feat(hygon): add Qwen3-235B-A3B BF16/W8A8 inference support

include/infinicore/adaptor/aten_adaptor.hpp

@@ -44,7 +44,7 @@ inline at::ScalarType to_at_dtype(DataType dtype) {
     case DataType::I64:
         return at::kLong;
     default:
-        throw std::runtime_error("Unsupported dtype for ATen");
+        throw std::runtime_error("Unsupported dtype for ATen: " + infinicore::toString(dtype));
     }
 }
 

include/infinicore/adaptor/lightop_adaptor.hpp

@@ -0,0 +1,147 @@
+#pragma once
+
+#include <cstddef>
+#include <string>
+
+namespace infinicore::adaptor::lightop {
+
+struct DeviceInfo {
+    std::string gpu_target;
+    int compute_units = 0;
+};
+
+DeviceInfo device_info(std::size_t device_index);
+
+} // namespace infinicore::adaptor::lightop
+
+#if defined(ENABLE_HYGON_API) && defined(ENABLE_ATEN)
+
+#include <ATen/ATen.h>
+
+#include <cstdint>
+#include <optional>
+
+namespace infinicore::adaptor::lightop {
+
+bool available();
+
+void preload_moe_w16a16_ops();
+
+void preload_moe_w16a16_ops(bool preload_legacy_gemm, bool preload_legacy_asm);
+
+void preload_moe_w16a16_marlin_asm(bool down_stage);
+
+void preload_moe_w8a8_ops();
+
+void preload_moe_align();
+
+void preload_moe_w8a8_marlin_asm();
+
+void preload_silu_and_mul();
+
+void preload_rms_rotary_embedding();
+
+void preload_reshape_and_cache_cuda();
+
+void fuse_silu_and_mul(
+    at::Tensor &input,
+    at::Tensor &output);
+
+void rms_rotary_embedding_fuse(
+    at::Tensor &positions,
+    at::Tensor &query,
+    at::Tensor &key,
+    int64_t head_size,
+    at::Tensor &cos_sin_cache,
+    bool is_neox,
+    at::Tensor q_weight,
+    at::Tensor k_weight,
+    const std::optional<at::Tensor> &q_bias = std::nullopt,
+    const std::optional<at::Tensor> &k_bias = std::nullopt,
+    double epsilon = 1e-6);
+
+void reshape_and_cache_cuda(
+    at::Tensor &key,
+    at::Tensor &value,
+    at::Tensor &key_cache,
+    at::Tensor &value_cache,
+    at::Tensor &slot_mapping,
+    const std::string &kv_cache_dtype,
+    at::Tensor &k_scale,
+    at::Tensor &v_scale);
+
+void moe_sum(
+    at::Tensor &input,
+    at::Tensor &output,
+    const std::optional<at::Tensor> &bias = std::nullopt,
+    const std::optional<at::Tensor> &expert_mask = std::nullopt,
+    const std::optional<at::Tensor> &local_num_tokens = std::nullopt,
+    float factor = 1.0f,
+    int expect_m = -1);
+
+void moe_align_block_size(
+    at::Tensor topk_ids,
+    int64_t num_experts,
+    int64_t block_size,
+    at::Tensor sorted_token_ids,
+    at::Tensor expert_ids,
+    at::Tensor num_tokens_post_padded,
+    const std::optional<at::Tensor> &expert_map = std::nullopt,
+    const std::optional<at::Tensor> &expert_mask = std::nullopt,
+    const std::optional<at::Tensor> &num_local_tokens = std::nullopt,
+    bool is_ep = false,
+    bool fuse_fill = true);
+void moe_gemm_marlin_w16a16(
+    at::Tensor input,
+    at::Tensor b_qweight,
+    at::Tensor output,
+    const std::optional<at::Tensor> &topk_weights,
+    at::Tensor sorted_token_ids,
+    at::Tensor expert_ids,
+    at::Tensor num_tokens_post_padded,
+    int64_t top_k,
+    int mode,
+    int delta);
+
+void moe_gemm_marlin_w8a8(
+    at::Tensor input,
+    at::Tensor b_qweight,
+    at::Tensor output,
+    at::Tensor a_scale,
+    at::Tensor b_scale,
+    const std::optional<at::Tensor> &topk_weights,
+    at::Tensor sorted_token_ids,
+    at::Tensor expert_ids,
+    at::Tensor num_tokens_post_padded,
+    int64_t top_k,
+    int mode,
+    int delta);
+
+void fuse_silu_mul_quant(
+    at::Tensor &input,
+    at::Tensor &output,
+    at::Tensor &scales,
+    std::optional<at::Tensor> &num_local_tokens,
+    int topk,
+    int expect_m,
+    std::optional<at::Tensor> &expert_ids);
+
+void preload_w8a8_linear_ops();
+
+void per_token_dynamic_quant_int8(
+    at::Tensor &output,
+    const at::Tensor &input,
+    at::Tensor &scales,
+    const at::Tensor &smooth);
+
+void blaslt_w8a8_gemm(
+    at::Tensor &output,
+    const at::Tensor &a,
+    const at::Tensor &b,
+    const at::Tensor &scale_a,
+    const at::Tensor &scale_b,
+    const std::optional<at::Tensor> &bias);
+
+} // namespace infinicore::adaptor::lightop
+
+#endif // ENABLE_HYGON_API && ENABLE_ATEN

include/infinicore/graph/graph.hpp

@@ -31,9 +31,9 @@ class DispatchableGraphOperator : public GraphOperator {
 protected:
     using run_schema = void (*)(void *);
     using cleanup_schema = void (*)(void **);
-    void *planned_meta_;
-    run_schema runner_;
-    cleanup_schema deleter_;
+    void *planned_meta_ = nullptr;
+    run_schema runner_ = nullptr;
+    cleanup_schema deleter_ = nullptr;
 };
 
 class Graph {

include/infinicore/nn/rope.hpp

@@ -85,6 +85,7 @@ class RoPE : public Module {
     double theta() const { return theta_; }
     Algo algo() const { return algo_; }
     DataType dtype() const { return dtype_; }
+    Tensor cos_sin_cache() const { return cos_sin_cache_; }
     const std::optional<std::vector<int>> &mrope_section() const { return mrope_section_; }
     bool mrope_interleaved() const { return mrope_interleaved_; }
     const Tensor &sin_cache() const { return sin_cache_; }
@@ -97,6 +98,7 @@ class RoPE : public Module {
     // Buffers (sin and cos cache tables) - not exposed in state_dict
     INFINICORE_NN_BUFFER(sin_cache);
     INFINICORE_NN_BUFFER(cos_cache);
+    INFINICORE_NN_BUFFER(cos_sin_cache);
 
 private:
     void initialize_cache();

include/infinicore/ops.hpp

@@ -44,6 +44,7 @@
 #include "ops/gelutanh.hpp"
 #include "ops/hardswish.hpp"
 #include "ops/hardtanh.hpp"
+#include "ops/hygon_moe_marlin.hpp"
 #include "ops/kv_caching.hpp"
 #include "ops/kimi_delta_attention.hpp"
 #include "ops/layer_norm.hpp"
@@ -54,14 +55,18 @@
 #include "ops/moe_align.hpp"
 #include "ops/moe_fused_dense.hpp"
 #include "ops/moe_fused_gate.hpp"
+#include "ops/moe_marlin_config.hpp"
 #include "ops/moe_sum.hpp"
 #include "ops/moe_topk_sigmoid.hpp"
 #include "ops/moe_topk_softmax.hpp"
+#include "ops/moe_w16a16_marlin.hpp"
+#include "ops/moe_w8a8_marlin.hpp"
 #include "ops/nrm2.hpp"
 #include "ops/ones.hpp"
 #include "ops/paged_attention.hpp"
 #include "ops/paged_attention_prefill.hpp"
 #include "ops/paged_caching.hpp"
+#include "ops/paged_flash_attention.hpp"
 #include "ops/per_tensor_dequant_i8.hpp"
 #include "ops/per_tensor_quant_i8.hpp"
 #include "ops/prepare_moe_input.hpp"
@@ -72,6 +77,7 @@
 #include "ops/recurrent_gated_delta_rule.hpp"
 #include "ops/relu.hpp"
 #include "ops/rms_norm.hpp"
+#include "ops/rms_rotary_embedding.hpp"
 #include "ops/rope.hpp"
 #include "ops/rot.hpp"
 #include "ops/rotg.hpp"
@@ -80,6 +86,7 @@
 #include "ops/rwkv5_wkv.hpp"
 #include "ops/scal.hpp"
 #include "ops/select_last_token_hidden.hpp"
+#include "ops/select_last_token_hidden_states.hpp"
 #include "ops/sigmoid.hpp"
 #include "ops/silu.hpp"
 #include "ops/silu_and_mul.hpp"

include/infinicore/ops/distributed/allreduce.hpp

@@ -12,6 +12,7 @@ class AllReduce : public graph::GraphOperator {
     AllReduce(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator);
     ~AllReduce();
     void run() const override;
+    bool is_device_graph_capture_safe() const override;
     static void execute(Tensor output, const Tensor &input, infinicclReduceOp_t op, infinicclComm_t communicator);
 
 private:

include/infinicore/ops/hygon_moe_marlin.hpp

@@ -0,0 +1,61 @@
+#pragma once
+
+#include "../tensor.hpp"
+
+#include <cstddef>
+
+namespace infinicore::op {
+
+enum class HygonMoeMarlinWeightFormat {
+    W16A16,
+    W8A8,
+};
+
+struct HygonMoeMarlinWeights {
+    Tensor packed_w13;
+    Tensor packed_w2;
+    Tensor packed_w13_scale;
+    Tensor packed_w2_scale;
+    HygonMoeMarlinWeightFormat format =
+        HygonMoeMarlinWeightFormat::W16A16;
+};
+
+struct HygonMoeMarlinWorkspace {
+    Tensor output;
+    Tensor cache13;
+    Tensor cache2;
+    Tensor input_i8;
+    Tensor input_scale;
+    Tensor cache2_i8;
+    Tensor cache2_scale;
+    Tensor sorted_token_ids;
+    Tensor expert_ids;
+    Tensor num_tokens_post_padded;
+
+    size_t cache13_capacity = 0;
+    size_t cache2_capacity = 0;
+    size_t sorted_token_ids_capacity = 0;
+    size_t expert_ids_capacity = 0;
+};
+
+struct HygonMoeMarlinOutput {
+    Tensor hidden_states;
+    Tensor sorted_token_ids;
+    Tensor expert_ids;
+    Tensor num_tokens_post_padded;
+    bool has_routing_metadata = false;
+};
+
+HygonMoeMarlinOutput hygon_moe_marlin_fused(
+    const Tensor &hidden_states,
+    const Tensor &topk_weights,
+    const Tensor &topk_ids,
+    const Tensor &expert_map,
+    const HygonMoeMarlinWeights &weights,
+    HygonMoeMarlinWorkspace &workspace,
+    size_t num_local_experts,
+    size_t hidden_size,
+    size_t intermediate_size,
+    size_t fallback_align_block_size);
+
+} // namespace infinicore::op

include/infinicore/ops/moe_marlin_config.hpp

@@ -0,0 +1,42 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../dtype.hpp"
+
+#include <cstddef>
+
+namespace infinicore::op {
+
+struct HygonMarlinGemmConfig {
+    int mode = 103;
+    int delta = 1;
+    size_t block_size_m = 16;
+    bool found = false;
+};
+
+struct HygonW16A16MarlinRuntimeConfig {
+    HygonMarlinGemmConfig gemm1;
+    HygonMarlinGemmConfig gemm2;
+    bool supported = false;
+};
+
+struct HygonW8A8MarlinRuntimeConfig {
+    HygonMarlinGemmConfig gemm1;
+    HygonMarlinGemmConfig gemm2;
+    bool supported = false;
+};
+
+HygonW16A16MarlinRuntimeConfig select_hygon_w16a16_marlin_config(
+    size_t num_tokens,
+    size_t hidden_size,
+    size_t intermediate_size,
+    DataType hidden_dtype,
+    size_t device_index);
+
+HygonW8A8MarlinRuntimeConfig select_hygon_w8a8_marlin_config(
+    size_t num_tokens,
+    size_t hidden_size,
+    size_t intermediate_size,
+    size_t device_index);
+
+} // namespace infinicore::op

include/infinicore/ops/moe_w16a16_marlin.hpp

@@ -0,0 +1,48 @@
+#pragma once
+
+#include "../device.hpp"
+#include "../graph/graph.hpp"
+#include "../tensor.hpp"
+#include "common/op.hpp"
+
+#include <cstddef>
+
+namespace infinicore::op {
+
+INFINICORE_GRAPH_OP_CLASS(MoeW16A16MarlinFusedDense,
+                          Tensor,
+                          Tensor,
+                          Tensor,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          const Tensor &,
+                          size_t,
+                          int,
+                          int,
+                          int,
+                          int);
+
+Tensor moe_w16a16_marlin_pack(const Tensor &weight);
+
+void moe_w16a16_marlin_fused_dense_(
+    Tensor output,
+    Tensor cache13,
+    Tensor cache2,
+    const Tensor &hidden_states,
+    const Tensor &w13_marlin,
+    const Tensor &w2_marlin,
+    const Tensor &topk_weights,
+    const Tensor &sorted_token_ids,
+    const Tensor &expert_ids,
+    const Tensor &num_tokens_post_padded,
+    size_t top_k,
+    int mode0,
+    int delta0,
+    int mode1,
+    int delta1);
+
+} // namespace infinicore::op