Fuse softmax_topk_oneshot with moe_compressed

riverlijunjie · riverlijunjie · commit 024511a6a64c · 2025-10-23T22:01:58.000+08:00
diff --git a/src/plugins/intel_gpu/include/intel_gpu/plugin/primitives_list.hpp b/src/plugins/intel_gpu/include/intel_gpu/plugin/primitives_list.hpp
@@ -311,4 +311,4 @@ REGISTER_FACTORY(internal, PagedAttentionExtension);
 REGISTER_FACTORY(internal, LoraSubgraph);
 REGISTER_FACTORY(internal, LoraSubgraphFused);
 REGISTER_FACTORY(internal, VLSDPA);
-REGISTER_FACTORY(internal, MOECompressed);
+REGISTER_FACTORY(internal, MOEFusedCompressed);
diff --git a/src/plugins/intel_gpu/include/intel_gpu/primitives/moe_fused_compressed.hpp b/src/plugins/intel_gpu/include/intel_gpu/primitives/moe_fused_compressed.hpp
@@ -6,45 +6,45 @@
 #include <vector>
 
 #include "intel_gpu/runtime/engine.hpp"
-#include "intel_gpu/op/moe_compressed.hpp"
+#include "intel_gpu/op/moe_fused_compressed.hpp"
 #include "primitive.hpp"
 
 namespace cldnn {
-using MOECompressed = ov::intel_gpu::op::MOECompressed;
+using MOEFusedCompressed = ov::intel_gpu::op::MOEFusedCompressed;
 
 /// @brief moe compressed primitive
 /// @details Performs moe compressed
-struct moe_compressed : public primitive_base<moe_compressed> {
-    CLDNN_DECLARE_PRIMITIVE(moe_compressed)
+struct moe_fused_compressed : public primitive_base<moe_fused_compressed> {
+    CLDNN_DECLARE_PRIMITIVE(moe_fused_compressed)
 
-    moe_compressed() : primitive_base("", {}) {}
+    moe_fused_compressed() : primitive_base("", {}) {}
 
     /// @brief Constructs moe primitive / layer.
     ///
     /// @param id                 An identifier of new primitive.
     /// @param inputs             A list of Input primitive ids (inputs).
-    moe_compressed(const primitive_id& id, const std::vector<input_info>& inputs, const MOECompressed::Config& config)
+    moe_fused_compressed(const primitive_id& id, const std::vector<input_info>& inputs, const MOEFusedCompressed::Config& config)
         : primitive_base(id, inputs, 1, {optional_data_type()}),
           _config(config) {}
 
-    MOECompressed::Config _config;
+    MOEFusedCompressed::Config _config;
 
     bool operator==(const primitive& rhs) const override {
         if (!compare_common_params(rhs))
             return false;
 
-        auto rhs_casted = downcast<const moe_compressed>(rhs);
+        auto rhs_casted = downcast<const moe_fused_compressed>(rhs);
 
         return std::memcmp(&_config, &rhs_casted._config, sizeof(_config)) == 0;
     }
 
     void save(BinaryOutputBuffer& ob) const override {
-        primitive_base<moe_compressed>::save(ob);
+        primitive_base<moe_fused_compressed>::save(ob);
         ob << make_data(&_config, sizeof(_config));
     }
 
     void load(BinaryInputBuffer& ib) override {
-        primitive_base<moe_compressed>::load(ib);
+        primitive_base<moe_fused_compressed>::load(ib);
         ib >> make_data(&_config, sizeof(_config));
     }
 };
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.cl b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.cl
@@ -92,9 +92,8 @@ KERNEL (gather_2d_ref)(
     #endif
 
     if (off == 0) {
-        // int top_idx = top_index[k];
-        // dst_rweight[k] = src_rweight[top_idx];
-        dst_rweight[k] = src_rweight[tok_idx];
+        int top_idx = top_index[k];
+        dst_rweight[k] = src_rweight[top_idx];
     }
 }
 
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.cpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.cpp
diff --git a/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.hpp b/src/plugins/intel_gpu/src/graph/impls/ocl_v2/moe_opt.hpp
@@ -22,16 +22,15 @@ namespace ov::intel_gpu::ocl {
 enum class MOEInputIndex : uint8_t {
     HIDDEN_STATES = 0,
     ROUTING_WEIGHTS = 1,
-    ROUTER_TOPK_OUTPUT_INDICES = 2,
-    WEIGHT_0 = 3,
-    SCALE_0 = 4,
-    ZP_0 = 5,
-    WEIGHT_1 = 6,
-    SCALE_1 = 7,
-    ZP_1 = 8,
-    WEIGHT_2 = 9,
-    SCALE_2 = 10,
-    ZP_2 = 11
+    WEIGHT_0 = 2,
+    SCALE_0 = 3,
+    ZP_0 = 4,
+    WEIGHT_1 = 5,
+    SCALE_1 = 6,
+    ZP_1 = 7,
+    WEIGHT_2 = 8,
+    SCALE_2 = 9,
+    ZP_2 = 10
 };
 
 struct MOEOpt : public ImplementationManager {
diff --git a/src/plugins/intel_gpu/src/graph/include/moe_inst.h b/src/plugins/intel_gpu/src/graph/include/moe_inst.h
@@ -8,21 +8,21 @@
 #include <string>
 #include <vector>
 
-#include "intel_gpu/primitives/moe_compressed.hpp"
+#include "intel_gpu/primitives/moe_fused_compressed.hpp"
 #include "primitive_inst.h"
 
 namespace cldnn {
 namespace details {}
 
 template <>
-struct typed_program_node<moe_compressed> : public typed_program_node_base<moe_compressed> {
+struct typed_program_node<moe_fused_compressed> : public typed_program_node_base<moe_fused_compressed> {
 private:
-    using parent = typed_program_node_base<moe_compressed>;
+    using parent = typed_program_node_base<moe_fused_compressed>;
 
 public:
     using parent::parent;
 
-    typed_program_node(std::shared_ptr<moe_compressed> prim, program& prog) : parent(prim, prog) {}
+    typed_program_node(std::shared_ptr<moe_fused_compressed> prim, program& prog) : parent(prim, prog) {}
 
     using parent::get_kernel_impl_params;
     std::unique_ptr<kernel_impl_params> get_kernel_impl_params(const std::vector<layout>& in_layouts, const std::vector<layout>& out_layouts) const override {
@@ -32,11 +32,11 @@ struct typed_program_node<moe_compressed> : public typed_program_node_base<moe_c
     }
 };
 
-using moe_node = typed_program_node<moe_compressed>;
+using moe_node = typed_program_node<moe_fused_compressed>;
 
 template <>
-class typed_primitive_inst<moe_compressed> : public typed_primitive_inst_base<moe_compressed> {
-    using parent = typed_primitive_inst_base<moe_compressed>;
+class typed_primitive_inst<moe_fused_compressed> : public typed_primitive_inst_base<moe_fused_compressed> {
+    using parent = typed_primitive_inst_base<moe_fused_compressed>;
     using parent::parent;
     using primitive_inst::update_output_memory;
 
@@ -48,5 +48,5 @@ class typed_primitive_inst<moe_compressed> : public typed_primitive_inst_base<mo
     typed_primitive_inst(network& network, const moe_node& node);
 };
 
-using moe_inst = typed_primitive_inst<moe_compressed>;
+using moe_inst = typed_primitive_inst<moe_fused_compressed>;
 }  // namespace cldnn
diff --git a/src/plugins/intel_gpu/src/graph/moe.cpp b/src/plugins/intel_gpu/src/graph/moe.cpp
@@ -13,7 +13,7 @@
 #include "program_node.h"
 
 namespace cldnn {
-GPU_DEFINE_PRIMITIVE_TYPE_ID(moe_compressed)
+GPU_DEFINE_PRIMITIVE_TYPE_ID(moe_fused_compressed)
 
 /*
     Calc_output_layout method is called only when output layout is invalidated.
diff --git a/src/plugins/intel_gpu/src/graph/registry/moe_impls.cpp b/src/plugins/intel_gpu/src/graph/registry/moe_impls.cpp
@@ -2,7 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-#include "intel_gpu/primitives/moe_compressed.hpp"
+#include "intel_gpu/primitives/moe_fused_compressed.hpp"
 #include "primitive_inst.h"
 #include "registry.hpp"
 
@@ -14,7 +14,7 @@ namespace ov::intel_gpu {
 
 using namespace cldnn;
 
-const std::vector<std::shared_ptr<cldnn::ImplementationManager>>& Registry<moe_compressed>::get_implementations() {
+const std::vector<std::shared_ptr<cldnn::ImplementationManager>>& Registry<moe_fused_compressed>::get_implementations() {
     static const std::vector<std::shared_ptr<ImplementationManager>> impls = {OV_GPU_CREATE_INSTANCE_OCL(ocl::MOEOpt, shape_types::any)};
 
     return impls;
diff --git a/src/plugins/intel_gpu/src/graph/registry/registry.hpp b/src/plugins/intel_gpu/src/graph/registry/registry.hpp
@@ -166,7 +166,7 @@ REGISTER_IMPLS(strided_slice);
 REGISTER_IMPLS(tile);
 REGISTER_IMPLS(col2im);
 REGISTER_IMPLS(vl_sdpa);
-REGISTER_IMPLS(moe_compressed);
+REGISTER_IMPLS(moe_fused_compressed);
 
 REGISTER_DEFAULT_IMPLS(assign, CPU_S, CPU_D);
 REGISTER_DEFAULT_IMPLS(read_value, CPU_S, CPU_D);
diff --git a/src/plugins/intel_gpu/src/plugin/ops/moe.cpp b/src/plugins/intel_gpu/src/plugin/ops/moe.cpp
@@ -2,35 +2,36 @@
 // SPDX-License-Identifier: Apache-2.0
 //
 
-#include "intel_gpu/op/moe_compressed.hpp"
+#include "intel_gpu/op/moe_fused_compressed.hpp"
 #include "intel_gpu/plugin/common_utils.hpp"
 #include "intel_gpu/plugin/program_builder.hpp"
-#include "intel_gpu/primitives/moe_compressed.hpp"
+#include "intel_gpu/primitives/moe_fused_compressed.hpp"
+#include "intel_gpu/primitives/moe_fused_compressed.hpp"
 
 
 namespace ov {
 namespace op {
 namespace internal {
-using MOECompressed = ov::intel_gpu::op::MOECompressed;
+using MOEFusedCompressed  = ov::intel_gpu::op::MOEFusedCompressed ;
 }  // namespace internal
 }  // namespace op
 }  // namespace ov
 
 namespace ov::intel_gpu {
 
-static void CreateMOECompressedOp(ProgramBuilder& p, const std::shared_ptr<ov::intel_gpu::op::MOECompressed>& op) {
+static void CreateMOEFusedCompressedOp(ProgramBuilder& p, const std::shared_ptr<ov::intel_gpu::op::MOEFusedCompressed>& op) {
     auto inputs = p.GetInputInfo(op);
     const auto& config = op->get_config();
-    OPENVINO_ASSERT(inputs.size() == 12, "Inputs count of MOECompressed should be 12");
+    OPENVINO_ASSERT(inputs.size() == 11, "Inputs count of MOEFusedCompressed should be 11");
 
     const std::string layerName = layer_type_name_ID(op);
     // auto& engine = p.get_engine();
 
-    const cldnn::moe_compressed moe(layerName, inputs, config);
+    const cldnn::moe_fused_compressed moe(layerName, inputs, config);
 
     p.add_primitive(*op, moe);
 }
 
-REGISTER_FACTORY_IMPL(internal, MOECompressed);
+REGISTER_FACTORY_IMPL(internal, MOEFusedCompressed);
 
 }  // namespace ov::intel_gpu
diff --git a/src/plugins/intel_gpu/src/plugin/transformations/convert_moe_to_compressed.cpp b/src/plugins/intel_gpu/src/plugin/transformations/convert_moe_to_compressed.cpp
@@ -117,7 +117,7 @@ ConvertMOEToMOECompressed::ConvertMOEToMOECompressed() {
         config.inter_size = weight_shape[1];
         config.num_expert = weight_shape[0];
         config.group_size = weight_shape[3];
-        config.top_k = topk_shape[1].get_length();
+        config.top_k = topk_shape[topk_shape.size() - 1].get_length();
         config.out_type = ov::element::f16;
         auto moe_compressed = std::make_shared<ov::intel_gpu::op::MOECompressed>(args, config);
 
@@ -144,7 +144,8 @@ ConvertMOEToMOECompressed::ConvertMOEToMOECompressed() {
         ov::copy_runtime_info(moe, moe_compressed);
         ov::replace_node(moe, moe_compressed);
 
-        std::cout << "ConvertMOEToMOECompressed is hit : config.top_k = " << config.top_k << std::endl;
+        std::cout << "ConvertMOEToMOECompressed is hit : num_expert = " << config.num_expert << ", top_k = " << config.top_k
+                  << ", hidden_size = " << config.hidden_size << ", inter_size = " << config.inter_size << ", group_size = " << config.group_size << std::endl;
         return true;
     };
 
diff --git a/src/plugins/intel_gpu/src/plugin/transformations/fuse_moe_compressed.cpp b/src/plugins/intel_gpu/src/plugin/transformations/fuse_moe_compressed.cpp
@@ -113,6 +113,9 @@ FuseMOECompressed::FuseMOECompressed() {
         ov::copy_runtime_info(moe_compressed, moe_fused_compressed);
         ov::replace_node(moe_compressed, moe_fused_compressed);
 
+        std::cout << "FuseMOECompressed is hit : num_expert = " << config.num_expert << ", top_k = " << config.top_k << ", hidden_size = " << config.hidden_size
+                  << ", inter_size = " << config.inter_size << ", group_size = " << config.group_size << std::endl;
+
         return true;
     };
 
diff --git a/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp b/src/plugins/intel_gpu/src/plugin/transformations_pipeline.cpp
@@ -81,6 +81,7 @@
 #include "plugin/transformations/fc_convert_fusion.hpp"
 #include "plugin/transformations/fc_horizontal_fusion.hpp"
 #include "plugin/transformations/fc_per_layer_scaling.hpp"
+#include "plugin/transformations/fuse_moe_compressed.hpp"
 #include "plugin/transformations/increase_position_ids_precision.hpp"
 #include "plugin/transformations/indirect_kv_cache.hpp"
 #include "plugin/transformations/kv_cache_compression.hpp"
@@ -400,6 +401,7 @@ void TransformationsPipeline::apply(std::shared_ptr<ov::Model> func) {
 
         manager.register_pass<ov::pass::FuseVectorizedMOE3GEMM>();
         manager.register_pass<ov::intel_gpu::ConvertMOEToMOECompressed>();
+        manager.register_pass<ov::intel_gpu::FuseMOECompressed>();
 
         manager.register_pass<ov::pass::InitNodeInfo>();
         manager.register_pass<EinsumDecomposition>();

Original file line number	Diff line number	Diff line change
`@@ -92,9 +92,8 @@ KERNEL (gather_2d_ref)(`
`92`	`92`	`#endif`
`93`	`93`
`94`	`94`	`if (off == 0) {`
`95`		`- // int top_idx = top_index[k];`
`96`		`- // dst_rweight[k] = src_rweight[top_idx];`
`97`		`- dst_rweight[k] = src_rweight[tok_idx];`
	`95`	`+ int top_idx = top_index[k];`
	`96`	`+ dst_rweight[k] = src_rweight[top_idx];`
`98`	`97`	`}`
`99`	`98`	`}`
`100`	`99`