diff --git a/.gitmodules b/.gitmodules index 08b00197e0..d0a88cb93e 100755 --- a/.gitmodules +++ b/.gitmodules @@ -20,7 +20,7 @@ fetchRecurseSubmodules = false [submodule "third_party/xllm_ops"] path = third_party/xllm_ops - url = https://gitcode.com/xLLM-AI/xllm_ops.git + url = https://github.com/xLLM-AI/xllm-ops.git fetchRecurseSubmodules = true [submodule "third_party/etcd_cpp_apiv3"] path = third_party/etcd_cpp_apiv3 @@ -36,7 +36,7 @@ fetchRecurseSubmodules = false [submodule "third_party/torch_npu_ops"] path = third_party/torch_npu_ops - url = https://gitcode.com/xLLM-AI/torch_npu_ops.git + url = https://github.com/xLLM-AI/torch-npu-ops.git fetchRecurseSubmodules = false [submodule "third_party/cutlass"] path = third_party/cutlass diff --git a/tests/core/framework/config/config_json_test.cpp b/tests/core/framework/config/config_json_test.cpp index f9e868b894..d989abc0b6 100644 --- a/tests/core/framework/config/config_json_test.cpp +++ b/tests/core/framework/config/config_json_test.cpp @@ -15,6 +15,7 @@ limitations under the License. #include +#include #include #include #include @@ -24,6 +25,7 @@ limitations under the License. #include "core/common/global_flags.h" #include "core/framework/config/config_utils.h" #include "core/framework/config/execution_config.h" +#include "core/framework/config/kernel_config.h" #include "core/framework/config/kv_cache_config.h" #include "core/framework/config/model_config.h" #include "core/framework/config/parallel_config.h" @@ -236,6 +238,38 @@ TEST(ConfigJsonTest, FromJsonUsesParsedOverrides) { EXPECT_EQ(scheduler_config.max_decode_token_per_sequence(), 256); } +#if defined(USE_NPU) +TEST(KernelConfigTest, MegaMoeDefaultsOffAndReadsJsonOverride) { + KernelConfig kernel_config; + EXPECT_EQ(kernel_config.mega_moe_mode(), "off"); + EXPECT_EQ(kernel_config.mega_moe_weight_cache_budget_bytes(), + 4LL * 1024 * 1024 * 1024); + + const JsonReader json = config::parse_json_string( + R"json({"mega_moe_mode": "on", "mega_moe_weight_cache_budget_bytes": 123456})json"); + kernel_config.from_json(json); + EXPECT_EQ(kernel_config.mega_moe_mode(), "on"); + EXPECT_EQ(kernel_config.mega_moe_weight_cache_budget_bytes(), 123456); + + nlohmann::ordered_json dumped_config; + kernel_config.append_config_json(dumped_config); + ASSERT_TRUE(dumped_config.contains("mega_moe_mode")); + EXPECT_EQ(dumped_config.at("mega_moe_mode").get(), "on"); + ASSERT_TRUE( + dumped_config.contains("mega_moe_weight_cache_budget_bytes")); + EXPECT_EQ(dumped_config.at("mega_moe_weight_cache_budget_bytes") + .get(), + 123456); + + const auto& names = KernelConfig::option_category().option_names; + EXPECT_NE(std::find(names.begin(), names.end(), "mega_moe_mode"), + names.end()); + EXPECT_NE(std::find(names.begin(), names.end(), + "mega_moe_weight_cache_budget_bytes"), + names.end()); +} +#endif + TEST(KVCacheConfigValidationTest, AcceptsSupportedIndexerCacheDtypes) { KVCacheConfig config; diff --git a/tests/core/framework/parallel_state/CMakeLists.txt b/tests/core/framework/parallel_state/CMakeLists.txt index 5e13dd3f92..efd3107e90 100644 --- a/tests/core/framework/parallel_state/CMakeLists.txt +++ b/tests/core/framework/parallel_state/CMakeLists.txt @@ -37,6 +37,19 @@ if(USE_NPU) hccl c_sec nnopbase) + + cc_test( + NAME + mega_moe_comm_resource_test + SRCS + mega_moe_comm_resource_test.cpp + DEPS + mega_moe_comm_resource + torch + GTest::gtest_main + hccl + hcomm + ) endif() if(USE_MLU) diff --git a/tests/core/framework/parallel_state/mega_moe_comm_resource_test.cpp b/tests/core/framework/parallel_state/mega_moe_comm_resource_test.cpp new file mode 100644 index 0000000000..ae674aa377 --- /dev/null +++ b/tests/core/framework/parallel_state/mega_moe_comm_resource_test.cpp @@ -0,0 +1,309 @@ +/* Copyright 2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include +#include +#include +#include +#include +#include + +#include "framework/parallel_state/mega_moe_comm_resource.h" + +namespace xllm { +namespace { + +MegaMoeCommSpec valid_spec() { + MegaMoeCommSpec spec; + spec.group_name = "ep_group"; + spec.hccl_comm = reinterpret_cast(0x1); + spec.ep_world_size = 16; + spec.device_index = 0; + spec.max_num_tokens_per_rank = 128; + return spec; +} + +TEST(MegaMoeCommResourceTest, AcceptsCompleteEp16Spec) { + const MegaMoeCommValidation validation = + validate_mega_moe_comm_spec(valid_spec()); + + EXPECT_TRUE(validation.valid); + EXPECT_EQ(validation.reason, MegaMoeCommRejectReason::NONE); +} + +TEST(MegaMoeCommResourceTest, AcceptsCompleteEp8Spec) { + MegaMoeCommSpec spec = valid_spec(); + spec.ep_world_size = 8; + + const MegaMoeCommValidation validation = + validate_mega_moe_comm_spec(spec); + + EXPECT_TRUE(validation.valid); + EXPECT_EQ(validation.reason, MegaMoeCommRejectReason::NONE); +} + +TEST(MegaMoeCommResourceTest, RejectsInvalidConstructionParameters) { + MegaMoeCommSpec spec = valid_spec(); + spec.group_name.clear(); + EXPECT_EQ(validate_mega_moe_comm_spec(spec).reason, + MegaMoeCommRejectReason::EMPTY_GROUP); + + spec = valid_spec(); + spec.hccl_comm = nullptr; + EXPECT_EQ(validate_mega_moe_comm_spec(spec).reason, + MegaMoeCommRejectReason::NULL_COMM); + + spec = valid_spec(); + spec.ep_world_size = 0; + EXPECT_EQ(validate_mega_moe_comm_spec(spec).reason, + MegaMoeCommRejectReason::UNSUPPORTED_EP_WORLD_SIZE); + + spec = valid_spec(); + spec.device_index = -1; + EXPECT_EQ(validate_mega_moe_comm_spec(spec).reason, + MegaMoeCommRejectReason::INVALID_DEVICE_INDEX); + + spec = valid_spec(); + spec.max_num_tokens_per_rank = 0; + EXPECT_EQ(validate_mega_moe_comm_spec(spec).reason, + MegaMoeCommRejectReason::INVALID_MAX_NUM_TOKENS_PER_RANK); +} + +TEST(MegaMoeCommResourceTest, FindsRequiredCann91KfcSymbols) { + const char* enabled = std::getenv("XLLM_RUN_CANN91_MEGA_MOE_SMOKE"); + if (enabled == nullptr || std::string(enabled) != "1") { + GTEST_SKIP() << "CANN 9.1 KFC symbol smoke is disabled; set " + "XLLM_RUN_CANN91_MEGA_MOE_SMOKE=1 in a CANN 9.1 " + "environment."; + } + const MegaMoeCommSymbolStatus status = probe_mega_moe_comm_symbols(); + + EXPECT_TRUE(status.available) << status.missing_symbol; + EXPECT_TRUE(status.missing_symbol.empty()); +} + +TEST(MegaMoeCommResourceTest, AcceptsAccessibleSpansEqualToLocalPayload) { + const auto validation = validate_mega_moe_buffer_accessible_spans( + 4096, std::vector(16, 4096)); + + EXPECT_TRUE(validation.valid); + EXPECT_EQ(validation.required_payload_size, 4096); +} + +TEST(MegaMoeCommResourceTest, + AcceptsLocalPayloadAndLargerRemoteIpcAccessibleSpans) { + constexpr uint64_t kLocalPayloadSize = 512ULL * 1024 * 1024; + constexpr uint64_t kObservedRemoteIpcOverhead = 1ULL * 1024 * 1024; + std::vector accessible_spans( + 16, kLocalPayloadSize + kObservedRemoteIpcOverhead); + accessible_spans[7] = kLocalPayloadSize; + + const auto validation = validate_mega_moe_buffer_accessible_spans( + kLocalPayloadSize, accessible_spans); + + EXPECT_TRUE(validation.valid); + EXPECT_EQ(validation.required_payload_size, kLocalPayloadSize); +} + +TEST(MegaMoeCommResourceTest, RejectsZeroLocalHcclPayloadSize) { + const auto validation = validate_mega_moe_buffer_accessible_spans( + 0, std::vector(16, 4096)); + + EXPECT_FALSE(validation.valid); + EXPECT_EQ(validation.required_payload_size, 0); +} + +TEST(MegaMoeCommResourceTest, + RejectsRemoteIpcAccessibleSpanSmallerThanLocalPayload) { + std::vector accessible_spans(16, 4096); + accessible_spans[7] = 2048; + + const auto validation = validate_mega_moe_buffer_accessible_spans( + 4096, accessible_spans); + + EXPECT_FALSE(validation.valid); + EXPECT_EQ(validation.mismatched_rank, 7); + EXPECT_EQ(validation.required_payload_size, 4096); + EXPECT_EQ(validation.accessible_span, 2048); +} + +TEST(MegaMoeCommResourceTest, + AcceptsDifferentRemoteIpcSpansThatCoverLocalPayload) { + std::vector accessible_spans(16, 4096); + accessible_spans[2] = 4097; + accessible_spans[9] = 64ULL * 1024 * 1024; + + const auto validation = validate_mega_moe_buffer_accessible_spans( + 4096, accessible_spans); + + EXPECT_TRUE(validation.valid); + EXPECT_EQ(validation.required_payload_size, 4096); +} + +TEST(MegaMoeCommResourceTest, RejectsZeroRemoteIpcAccessibleSpan) { + std::vector accessible_spans(16, 4096); + accessible_spans[11] = 0; + + const auto validation = validate_mega_moe_buffer_accessible_spans( + 4096, accessible_spans); + + EXPECT_FALSE(validation.valid); + EXPECT_EQ(validation.mismatched_rank, 11); + EXPECT_EQ(validation.required_payload_size, 4096); + EXPECT_EQ(validation.accessible_span, 0); +} + +std::shared_ptr fake_resource(int value) { + auto owner = std::make_shared(value); + return std::shared_ptr( + owner, reinterpret_cast(owner.get())); +} + +TEST(MegaMoeCommResourceTest, SharedSlotCreatesOnceForSameCompleteKey) { + MegaMoeCommResourceSlot slot; + int create_count = 0; + const auto factory = [&](const MegaMoeCommSpec&) { + return fake_resource(++create_count); + }; + + auto first = slot.acquire(valid_spec(), factory); + auto second = slot.acquire(valid_spec(), factory); + + EXPECT_EQ(create_count, 1); + EXPECT_EQ(first.get(), second.get()); +} + +TEST(MegaMoeCommResourceTest, SharedSlotCreatesNewResourceWhenKeyChanges) { + MegaMoeCommResourceSlot slot; + int create_count = 0; + const auto factory = [&](const MegaMoeCommSpec&) { + return fake_resource(++create_count); + }; + auto spec = valid_spec(); + auto previous = slot.acquire(spec, factory); + + spec.hccl_comm = reinterpret_cast(0x2); + auto changed_comm = slot.acquire(spec, factory); + EXPECT_NE(previous.get(), changed_comm.get()); + previous = changed_comm; + + spec.group_name = "other_ep_group"; + auto changed_group = slot.acquire(spec, factory); + EXPECT_NE(previous.get(), changed_group.get()); + previous = changed_group; + + spec.ep_world_size = 32; + auto changed_world = slot.acquire(spec, factory); + EXPECT_NE(previous.get(), changed_world.get()); + previous = changed_world; + + spec.device_index = 1; + auto changed_device = slot.acquire(spec, factory); + EXPECT_NE(previous.get(), changed_device.get()); + previous = changed_device; + + spec.max_num_tokens_per_rank = 256; + auto changed_envelope = slot.acquire(spec, factory); + EXPECT_NE(previous.get(), changed_envelope.get()); + EXPECT_EQ(create_count, 6); +} + +TEST(MegaMoeCommResourceTest, SharedSlotDoesNotCacheFactoryFailure) { + MegaMoeCommResourceSlot slot; + int create_count = 0; + EXPECT_THROW( + slot.acquire(valid_spec(), [&](const MegaMoeCommSpec&) { + ++create_count; + throw std::runtime_error("injected create failure"); + return std::shared_ptr(); + }), + std::runtime_error); + + auto resource = slot.acquire(valid_spec(), [&](const MegaMoeCommSpec&) { + return fake_resource(++create_count); + }); + EXPECT_NE(resource, nullptr); + EXPECT_EQ(create_count, 2); +} + +TEST(MegaMoeCommResourceTest, ExplicitResetReleasesBeforeCommunicatorShutdown) { + std::vector events; + MegaMoeCommResourceSlot slot; + auto injected = std::shared_ptr( + reinterpret_cast(0x1), + [&](MegaMoeCommResource*) { events.push_back("resource_released"); }); + auto layer_reference = slot.acquire( + valid_spec(), + [&](const MegaMoeCommSpec&) { return injected; }); + + injected.reset(); + layer_reference.reset(); + EXPECT_TRUE(events.empty()); + + // ProcessGroupImpl performs this explicit reset before shutdown/destroy of + // its HCCL communicator. + slot.reset(); + events.push_back("communicator_shutdown"); + + ASSERT_EQ(events.size(), 2); + EXPECT_EQ(events[0], "resource_released"); + EXPECT_EQ(events[1], "communicator_shutdown"); +} + +TEST(MegaMoeCommResourceTest, LayerWeakReferenceDoesNotExtendSlotLifetime) { + MegaMoeCommResourceSlot slot; + auto acquired = slot.acquire( + valid_spec(), + [&](const MegaMoeCommSpec&) { return fake_resource(1); }); + std::weak_ptr layer_reference = acquired; + + EXPECT_EQ(acquired.use_count(), 2); + acquired.reset(); + EXPECT_FALSE(layer_reference.expired()); + + slot.reset_for_teardown(); + EXPECT_TRUE(layer_reference.expired()); +} + +TEST(MegaMoeCommResourceTest, ActiveForwardLockBlocksTeardownGuard) { + MegaMoeCommResourceSlot slot; + auto acquired = slot.acquire( + valid_spec(), + [&](const MegaMoeCommSpec&) { return fake_resource(1); }); + std::weak_ptr layer_reference = acquired; + acquired.reset(); + auto active_forward_lock = layer_reference.lock(); + ASSERT_NE(active_forward_lock, nullptr); + + EXPECT_DEATH(slot.reset_for_teardown(), + "active MegaMoe forward"); + + active_forward_lock.reset(); + slot.reset_for_teardown(); + EXPECT_TRUE(layer_reference.expired()); +} + +TEST(MegaMoeCommResourceTest, EnforcesSingleOwnerLifecycle) { + EXPECT_TRUE(std::is_destructible_v); + EXPECT_FALSE(std::is_copy_constructible_v); + EXPECT_FALSE(std::is_copy_assignable_v); + EXPECT_FALSE(std::is_move_constructible_v); + EXPECT_FALSE(std::is_move_assignable_v); +} + +} // namespace +} // namespace xllm diff --git a/tests/core/kernels/npu/CMakeLists.txt b/tests/core/kernels/npu/CMakeLists.txt index 97ad595459..ce5452d049 100644 --- a/tests/core/kernels/npu/CMakeLists.txt +++ b/tests/core/kernels/npu/CMakeLists.txt @@ -1,5 +1,19 @@ include(cc_test) +cc_test( + NAME + mega_moe_test + SRCS + mega_moe_test.cpp + DEPS + ascendcl + torch + torch_npu + npu_kernels + GTest::gtest_main + glog::glog +) + # Temporarily disabled: Dsv4ScatterCache tests are unstable in the current # NPU test environment. # cc_test( diff --git a/tests/core/kernels/npu/mega_moe_test.cpp b/tests/core/kernels/npu/mega_moe_test.cpp new file mode 100644 index 0000000000..2c4bbf375c --- /dev/null +++ b/tests/core/kernels/npu/mega_moe_test.cpp @@ -0,0 +1,147 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include +#include +#include + +#include "core/kernels/npu/mega_moe_acl_contract.h" +#include "core/kernels/npu/npu_ops_api.h" + +namespace xllm::kernel::npu { +namespace { + +TEST(MegaMoeTest, AcceptsOnlyExactCustomVendorAbiProvenance) { + const std::string expected = + "vendor/op_api/lib/libcust_opapi.so"; + auto compatible = validate_mega_moe_op_api_paths( + expected, expected, expected); + EXPECT_TRUE(compatible.compatible); + EXPECT_TRUE(compatible.same_library); + + auto stock = validate_mega_moe_op_api_paths( + expected, + "stock/op_api/lib/libopapi.so", + "stock/op_api/lib/libopapi.so"); + EXPECT_FALSE(stock.compatible); + EXPECT_TRUE(stock.same_library); + + auto split = validate_mega_moe_op_api_paths( + expected, expected, "other/vendor/libcust_opapi.so"); + EXPECT_FALSE(split.compatible); + EXPECT_FALSE(split.same_library); +} + +TEST(MegaMoeTest, CustomAbiKeepsOptionalInputsNullAndTopologyBeforeOutputs) { + const auto bf16 = torch::TensorOptions().dtype(torch::kBFloat16); + const auto int32 = torch::TensorOptions().dtype(torch::kInt32); + auto context = torch::empty({1}, int32); + auto x = torch::empty({1, 2048}, bf16); + auto ids = torch::empty({1, 8}, int32); + auto weights = torch::empty({1, 8}, bf16); + std::vector w1_storage = { + torch::empty({2048, 1024}, bf16)}; + std::vector w2_storage = { + torch::empty({512, 2048}, bf16)}; + torch::TensorList w1(w1_storage); + torch::TensorList w2(w2_storage); + std::optional missing_list = std::nullopt; + std::optional missing_mask = std::nullopt; + auto output = torch::empty_like(x); + auto token_counts = torch::empty({1}, int32); + int64_t moe_expert_num = 256; + int64_t ep_world_size = 16; + int64_t ccl_buffer_size = 1; + int64_t max_recv_token_num = 0; + int64_t dispatch_quant_mode = 0; + int64_t dispatch_quant_out_dtype = 28; + int64_t combine_quant_mode = 0; + char comm_alg_storage[] = ""; + char* comm_alg = comm_alg_storage; + int64_t num_max_tokens_per_rank = 1; + char activation_storage[] = "swiglu"; + char* activation = activation_storage; + float activation_clamp = 3.4e38F; + int64_t topo_type = 0; + int64_t rank_num_per_server = 2; + + bool called = false; + auto fake_workspace = [&](auto&... arguments) { + auto args = std::forward_as_tuple(arguments...); + static_assert(std::tuple_size_v == 26); + EXPECT_FALSE(std::get<6>(args).has_value()); + EXPECT_FALSE(std::get<7>(args).has_value()); + EXPECT_FALSE(std::get<8>(args).has_value()); + EXPECT_FALSE(std::get<9>(args).has_value()); + EXPECT_FALSE(std::get<10>(args).has_value()); + EXPECT_EQ(std::get<22>(args), 0); + EXPECT_EQ(std::get<23>(args), 2); + EXPECT_EQ(&std::get<24>(args), &output); + EXPECT_EQ(&std::get<25>(args), &token_counts); + called = true; + }; + execute_mega_moe_acl_contract(fake_workspace, + context, x, ids, weights, w1, w2, + missing_list, missing_list, + missing_list, missing_list, missing_mask, + moe_expert_num, ep_world_size, ccl_buffer_size, + max_recv_token_num, dispatch_quant_mode, + dispatch_quant_out_dtype, combine_quant_mode, + comm_alg, num_max_tokens_per_rank, activation, + activation_clamp, topo_type, + rank_num_per_server, output, token_counts); + EXPECT_TRUE(called); +} + +TEST(MegaMoeTest, ExactVendorExportsBothAclnnSymbols) { + const char* enabled = std::getenv("XLLM_RUN_CANN91_MEGA_MOE_SMOKE"); + if (enabled == nullptr || std::string(enabled) != "1") { + GTEST_SKIP() << "CANN 9.1 custom-vendor ABI smoke is disabled; set " + "XLLM_RUN_CANN91_MEGA_MOE_SMOKE=1 in the isolated vendor " + "environment."; + } + EXPECT_TRUE(has_mega_moe()); +} + +TEST(MegaMoeTest, RejectsNonMatrixInputBeforeLaunch) { + const auto bf16 = torch::TensorOptions().dtype(torch::kBFloat16); + const auto int32 = torch::TensorOptions().dtype(torch::kInt); + const auto fp32 = torch::TensorOptions().dtype(torch::kFloat); + const torch::Tensor context = torch::empty({1}, int32); + const torch::Tensor x = torch::empty({2048}, bf16); + const torch::Tensor topk_ids = torch::empty({1, 8}, int32); + const torch::Tensor topk_weights = torch::empty({1, 8}, fp32); + const std::vector weight1 = { + torch::empty({2048, 1024}, bf16)}; + const std::vector weight2 = { + torch::empty({512, 2048}, bf16)}; + + EXPECT_THROW(apply_npu_mega_moe(context, + x, + topk_ids, + topk_weights, + weight1, + weight2, + 256, + 16, + 1), + c10::Error); +} + +} // namespace +} // namespace xllm::kernel::npu diff --git a/tests/core/layers/npu_torch/CMakeLists.txt b/tests/core/layers/npu_torch/CMakeLists.txt index 09162f4f2e..77540014f6 100644 --- a/tests/core/layers/npu_torch/CMakeLists.txt +++ b/tests/core/layers/npu_torch/CMakeLists.txt @@ -1,5 +1,26 @@ include(cc_test) +cc_test( + NAME + mega_moe_policy_test + SRCS + mega_moe_policy_test.cpp + DEPS + :mega_moe_policy + GTest::gtest_main +) + +cc_test( + NAME + mega_moe_runtime_test + SRCS + mega_moe_runtime_test.cpp + DEPS + :mega_moe_runtime + torch + GTest::gtest_main +) + cc_test( NAME npu_deepseek_v4_indexer_test diff --git a/tests/core/layers/npu_torch/mega_moe_policy_test.cpp b/tests/core/layers/npu_torch/mega_moe_policy_test.cpp new file mode 100644 index 0000000000..a4b53916a8 --- /dev/null +++ b/tests/core/layers/npu_torch/mega_moe_policy_test.cpp @@ -0,0 +1,261 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include + +#include +#include + +#include "layers/npu_torch/mega_moe_policy.h" + +namespace xllm { +namespace layer { +namespace { + +MegaMoeCapability supported_capability() { + MegaMoeCapability capability; + capability.soc_name = "Ascend910_93"; + capability.cann_version = "9.1.0-beta.3"; + capability.vendor_version = "9.1.0-beta.3"; + capability.dtype = MegaMoeDType::BFLOAT16; + capability.activation = "swiglu"; + capability.hidden_size = 2048; + capability.intermediate_size = 512; + capability.num_experts = 256; + capability.topk = 8; + capability.ep_size = 16; + capability.tp_size = 1; + capability.dp_size = 16; + capability.workspace_symbol_ready = true; + capability.execute_symbol_ready = true; + capability.comm_context_ready = true; + return capability; +} + +void expect_rejected(const MegaMoeCapability& capability, + MegaMoeRejectionReason reason) { + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::ON, capability); + EXPECT_FALSE(decision.use_mega_moe); + EXPECT_TRUE(decision.fatal); + EXPECT_EQ(decision.reason, reason); +} + +TEST(MegaMoePolicyTest, ParsesExplicitModes) { + EXPECT_EQ(parse_mega_moe_mode("off"), MegaMoeMode::OFF); + EXPECT_EQ(parse_mega_moe_mode("auto"), MegaMoeMode::AUTO); + EXPECT_EQ(parse_mega_moe_mode("on"), MegaMoeMode::ON); + EXPECT_EQ(parse_mega_moe_mode("invalid"), std::nullopt); +} + +TEST(MegaMoePolicyTest, OffKeepsLegacyExecution) { + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::OFF, supported_capability()); + EXPECT_FALSE(decision.use_mega_moe); + EXPECT_FALSE(decision.fatal); + EXPECT_EQ(decision.reason, MegaMoeRejectionReason::DISABLED); +} + +TEST(MegaMoePolicyTest, OnAndAutoSelectExactWhitelist) { + for (MegaMoeMode mode : {MegaMoeMode::ON, MegaMoeMode::AUTO}) { + const MegaMoeDecision decision = + decide_mega_moe(mode, supported_capability()); + EXPECT_TRUE(decision.use_mega_moe); + EXPECT_FALSE(decision.fatal); + EXPECT_EQ(decision.reason, MegaMoeRejectionReason::NONE); + } +} + +TEST(MegaMoePolicyTest, OnSupportsEp16MoeTp1WithAttentionDp16) { + MegaMoeCapability capability = supported_capability(); + capability.ep_size = 16; + capability.tp_size = 1; + capability.dp_size = 16; + + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::ON, capability); + EXPECT_TRUE(decision.use_mega_moe); + EXPECT_FALSE(decision.fatal); + EXPECT_EQ(decision.reason, MegaMoeRejectionReason::NONE); +} + +TEST(MegaMoePolicyTest, OnSupportsEp8MoeTp1WithAttentionDp8) { + MegaMoeCapability capability = supported_capability(); + capability.ep_size = 8; + capability.tp_size = 1; + capability.dp_size = 8; + + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::ON, capability); + + EXPECT_TRUE(decision.use_mega_moe); + EXPECT_FALSE(decision.fatal); + EXPECT_EQ(decision.reason, MegaMoeRejectionReason::NONE); +} + +TEST(MegaMoePolicyTest, AttentionDpDoesNotAffectOperatorEligibility) { + for (int64_t dp_size : {1, 2, 8, 16, 32}) { + SCOPED_TRACE(dp_size); + MegaMoeCapability capability = supported_capability(); + capability.dp_size = dp_size; + + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::ON, capability); + EXPECT_TRUE(decision.use_mega_moe); + EXPECT_FALSE(decision.fatal); + EXPECT_EQ(decision.reason, MegaMoeRejectionReason::NONE); + } +} + +TEST(MegaMoePolicyTest, AutoFallsBackOnlyBeforeCollectivesStart) { + MegaMoeCapability capability = supported_capability(); + capability.ep_size = 3; + const MegaMoeDecision before_collectives = + decide_mega_moe(MegaMoeMode::AUTO, capability); + EXPECT_FALSE(before_collectives.use_mega_moe); + EXPECT_FALSE(before_collectives.fatal); + EXPECT_EQ(before_collectives.reason, MegaMoeRejectionReason::UNSUPPORTED_EP); + + capability.collectives_started = true; + const MegaMoeDecision after_collectives = + decide_mega_moe(MegaMoeMode::AUTO, capability); + EXPECT_FALSE(after_collectives.use_mega_moe); + EXPECT_TRUE(after_collectives.fatal); + EXPECT_EQ(after_collectives.reason, + MegaMoeRejectionReason::COLLECTIVES_ALREADY_STARTED); +} + +TEST(MegaMoePolicyTest, WeightBudgetIsFatalOnAndFallbackAuto) { + MegaMoeCapability capability = supported_capability(); + capability.weight_budget_ready = false; + + const auto forced = decide_mega_moe(MegaMoeMode::ON, capability); + EXPECT_FALSE(forced.use_mega_moe); + EXPECT_TRUE(forced.fatal); + EXPECT_EQ(forced.reason, + MegaMoeRejectionReason::WEIGHT_BUDGET_EXCEEDED); + + const auto automatic = decide_mega_moe(MegaMoeMode::AUTO, capability); + EXPECT_FALSE(automatic.use_mega_moe); + EXPECT_FALSE(automatic.fatal); + EXPECT_EQ(automatic.reason, + MegaMoeRejectionReason::WEIGHT_BUDGET_EXCEEDED); +} + +TEST(MegaMoePolicyTest, FusedMc2ConflictFailsFast) { + MegaMoeCapability capability = supported_capability(); + capability.fused_mc2_enabled = true; + expect_rejected(capability, MegaMoeRejectionReason::FUSED_MC2_CONFLICT); + + const MegaMoeDecision auto_decision = + decide_mega_moe(MegaMoeMode::AUTO, capability); + EXPECT_FALSE(auto_decision.use_mega_moe); + EXPECT_TRUE(auto_decision.fatal); + EXPECT_EQ(auto_decision.reason, + MegaMoeRejectionReason::FUSED_MC2_CONFLICT); +} + +TEST(MegaMoePolicyTest, RejectsUnsupportedPlatformAndVersions) { + MegaMoeCapability capability = supported_capability(); + capability.soc_name = "Ascend910B"; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_SOC); + capability = supported_capability(); + capability.cann_version = "9.0.0"; + expect_rejected(capability, + MegaMoeRejectionReason::UNSUPPORTED_CANN_VERSION); + capability = supported_capability(); + capability.vendor_version = "9.0.0"; + expect_rejected(capability, + MegaMoeRejectionReason::UNSUPPORTED_VENDOR_VERSION); +} + +TEST(MegaMoePolicyTest, RejectsMissingSymbolsAndContext) { + MegaMoeCapability capability = supported_capability(); + capability.workspace_symbol_ready = false; + expect_rejected(capability, MegaMoeRejectionReason::OP_SYMBOLS_UNAVAILABLE); + capability = supported_capability(); + capability.execute_symbol_ready = false; + expect_rejected(capability, MegaMoeRejectionReason::OP_SYMBOLS_UNAVAILABLE); + capability = supported_capability(); + capability.comm_context_ready = false; + expect_rejected(capability, + MegaMoeRejectionReason::COMM_CONTEXT_UNAVAILABLE); +} + +TEST(MegaMoePolicyTest, RejectsUnsupportedDtypeActivationAndShape) { + MegaMoeCapability capability = supported_capability(); + capability.dtype = MegaMoeDType::OTHER; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_DTYPE); + capability = supported_capability(); + capability.activation = "gelu"; + expect_rejected(capability, + MegaMoeRejectionReason::UNSUPPORTED_ACTIVATION); + capability = supported_capability(); + capability.hidden_size = 4096; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_SHAPE); + capability = supported_capability(); + capability.intermediate_size = 1024; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_SHAPE); + capability = supported_capability(); + capability.num_experts = 128; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_SHAPE); + capability = supported_capability(); + capability.topk = 4; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_SHAPE); +} + +TEST(MegaMoePolicyTest, RejectsUnsupportedParallelAndRuntimeModes) { + MegaMoeCapability capability = supported_capability(); + capability.ep_size = 3; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_EP); + capability = supported_capability(); + capability.tp_size = 2; + expect_rejected(capability, MegaMoeRejectionReason::UNSUPPORTED_TP); + capability = supported_capability(); + capability.graph_enabled = true; + expect_rejected(capability, MegaMoeRejectionReason::GRAPH_ENABLED); + capability = supported_capability(); + capability.quantization_enabled = true; + expect_rejected(capability, + MegaMoeRejectionReason::QUANTIZATION_ENABLED); +} + +TEST(MegaMoePolicyTest, MegaMoePreservesGlobalRoutingAndAddsSharedOnce) { + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::ON, supported_capability()); + const MegaMoeExecutionContract contract = + make_mega_moe_execution_contract( + decision, /*ep_size=*/16, /*has_shared_expert=*/true); + EXPECT_TRUE(contract.preserve_global_topk); + EXPECT_FALSE(contract.apply_local_expert_mask); + EXPECT_FALSE(contract.reduce_routed_output_across_ep); + EXPECT_EQ(contract.shared_output_additions, 1); +} + +TEST(MegaMoePolicyTest, LegacyMasksLocalRoutingAndReducesAcrossEp) { + const MegaMoeDecision decision = + decide_mega_moe(MegaMoeMode::OFF, supported_capability()); + const MegaMoeExecutionContract contract = + make_mega_moe_execution_contract( + decision, /*ep_size=*/16, /*has_shared_expert=*/true); + EXPECT_FALSE(contract.preserve_global_topk); + EXPECT_TRUE(contract.apply_local_expert_mask); + EXPECT_TRUE(contract.reduce_routed_output_across_ep); + EXPECT_EQ(contract.shared_output_additions, 1); +} + +} // namespace +} // namespace layer +} // namespace xllm diff --git a/tests/core/layers/npu_torch/mega_moe_runtime_test.cpp b/tests/core/layers/npu_torch/mega_moe_runtime_test.cpp new file mode 100644 index 0000000000..71b080ce04 --- /dev/null +++ b/tests/core/layers/npu_torch/mega_moe_runtime_test.cpp @@ -0,0 +1,285 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include + +#include +#include + +#include "layers/npu_torch/mega_moe_runtime.h" + +namespace xllm::layer { +namespace { + +TEST(MegaMoeRuntimeTest, DefersCacheUntilW13ThenW2ShardsAreComplete) { + int build_count = 0; + const auto observe = [&](bool ready, bool w13_loaded, bool w2_loaded) { + const auto action = plan_mega_moe_weight_cache( + true, ready, w13_loaded, w2_loaded, false); + if (action == MegaMoeWeightCacheAction::BUILD) { + ++build_count; + } + return action; + }; + + EXPECT_EQ(observe(false, true, false), + MegaMoeWeightCacheAction::WAIT_FOR_WEIGHTS); + EXPECT_EQ(build_count, 0); + EXPECT_EQ(observe(false, true, true), + MegaMoeWeightCacheAction::BUILD); + EXPECT_EQ(build_count, 1); + EXPECT_EQ(observe(true, true, true), + MegaMoeWeightCacheAction::SKIP); + EXPECT_EQ(build_count, 1); +} + +TEST(MegaMoeRuntimeTest, DefersCacheUntilW2ThenW13ShardsAreComplete) { + int build_count = 0; + const auto observe = [&](bool ready, bool w13_loaded, bool w2_loaded) { + const auto action = plan_mega_moe_weight_cache( + true, ready, w13_loaded, w2_loaded, false); + if (action == MegaMoeWeightCacheAction::BUILD) { + ++build_count; + } + return action; + }; + + EXPECT_EQ(observe(false, false, true), + MegaMoeWeightCacheAction::WAIT_FOR_WEIGHTS); + EXPECT_EQ(build_count, 0); + EXPECT_EQ(observe(false, true, true), + MegaMoeWeightCacheAction::BUILD); + EXPECT_EQ(build_count, 1); + EXPECT_EQ(observe(true, true, true), + MegaMoeWeightCacheAction::SKIP); + EXPECT_EQ(build_count, 1); +} + +TEST(MegaMoeRuntimeTest, ForcedForwardFailsClosedWhenShardIsMissing) { + EXPECT_EQ(plan_mega_moe_weight_cache(true, false, true, false, true), + MegaMoeWeightCacheAction::FAIL_MISSING_WEIGHTS); + EXPECT_EQ(plan_mega_moe_weight_cache(true, false, false, true, true), + MegaMoeWeightCacheAction::FAIL_MISSING_WEIGHTS); +} + +TEST(MegaMoeRuntimeTest, DisabledModesNeverBuildOrRequireWeightCache) { + EXPECT_EQ(plan_mega_moe_weight_cache(false, false, true, true, false), + MegaMoeWeightCacheAction::SKIP); + EXPECT_EQ(plan_mega_moe_weight_cache(false, false, false, false, true), + MegaMoeWeightCacheAction::SKIP); +} + +TEST(MegaMoeRuntimeTest, BuildsIndependentContiguousWeightCaches) { + constexpr int64_t kExperts = 3; + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + torch::Tensor canonical_w13 = + torch::arange(kExperts * 2 * kIntermediate * kHidden, options) + .reshape({kExperts, 2 * kIntermediate, kHidden}); + torch::Tensor canonical_w2 = + torch::arange(kExperts * kHidden * kIntermediate, options) + .reshape({kExperts, kHidden, kIntermediate}); + const torch::Tensor original_w13 = canonical_w13.clone(); + const torch::Tensor original_w2 = canonical_w2.clone(); + + MegaMoeWeightCache cache = build_mega_moe_weight_cache( + canonical_w13, canonical_w2, kHidden, kIntermediate); + + ASSERT_TRUE(cache.ready()); + EXPECT_EQ(cache.weight1_storage.sizes(), + torch::IntArrayRef({kExperts, kHidden, 2 * kIntermediate})); + EXPECT_EQ(cache.weight2_storage.sizes(), + torch::IntArrayRef({kExperts, kIntermediate, kHidden})); + EXPECT_TRUE(cache.weight1_storage.is_contiguous()); + EXPECT_TRUE(cache.weight2_storage.is_contiguous()); + EXPECT_NE(cache.weight1_storage.data_ptr(), canonical_w13.data_ptr()); + EXPECT_NE(cache.weight2_storage.data_ptr(), canonical_w2.data_ptr()); + EXPECT_EQ(cache.weight1.size(), kExperts); + EXPECT_EQ(cache.weight2.size(), kExperts); + EXPECT_TRUE(torch::equal(cache.weight1[1], canonical_w13[1].transpose(0, 1))); + EXPECT_TRUE(torch::equal(cache.weight2[2], canonical_w2[2].transpose(0, 1))); + + cache.weight1_storage.zero_(); + cache.weight2_storage.zero_(); + EXPECT_TRUE(torch::equal(canonical_w13, original_w13)); + EXPECT_TRUE(torch::equal(canonical_w2, original_w2)); + + const int64_t expected_bytes = + (canonical_w13.numel() + canonical_w2.numel()) * + static_cast(canonical_w13.element_size()); + EXPECT_EQ(cache.memory_bytes, expected_bytes); +} + +TEST(MegaMoeRuntimeTest, AccountsSuccessfulWeightCacheAllocation) { + constexpr int64_t kExperts = 2; + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + const auto canonical_w13 = + torch::empty({kExperts, 2 * kIntermediate, kHidden}, options); + const auto canonical_w2 = + torch::empty({kExperts, kHidden, kIntermediate}, options); + const int64_t expected_bytes = + (canonical_w13.numel() + canonical_w2.numel()) * + static_cast(canonical_w13.element_size()); + const int64_t before = current_mega_moe_weight_cache_bytes(); + MegaMoeWeightBudgetStatus status; + + auto cache = try_build_mega_moe_weight_cache(canonical_w13, + canonical_w2, + kHidden, + kIntermediate, + expected_bytes, + &status); + + ASSERT_TRUE(cache.has_value()); + EXPECT_TRUE(status.allowed); + EXPECT_EQ(status.estimated_bytes, expected_bytes); + EXPECT_EQ(status.current_bytes, before); + EXPECT_EQ(status.limit_bytes, expected_bytes); + EXPECT_EQ(current_mega_moe_weight_cache_bytes(), before + expected_bytes); +} + +TEST(MegaMoeRuntimeTest, RejectsWeightCacheAllocationOverBudget) { + constexpr int64_t kExperts = 2; + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + const auto canonical_w13 = + torch::empty({kExperts, 2 * kIntermediate, kHidden}, options); + const auto canonical_w2 = + torch::empty({kExperts, kHidden, kIntermediate}, options); + const int64_t expected_bytes = + (canonical_w13.numel() + canonical_w2.numel()) * + static_cast(canonical_w13.element_size()); + const int64_t before = current_mega_moe_weight_cache_bytes(); + MegaMoeWeightBudgetStatus status; + + auto cache = try_build_mega_moe_weight_cache(canonical_w13, + canonical_w2, + kHidden, + kIntermediate, + expected_bytes - 1, + &status); + + EXPECT_FALSE(cache.has_value()); + EXPECT_FALSE(status.allowed); + EXPECT_EQ(status.estimated_bytes, expected_bytes); + EXPECT_EQ(status.current_bytes, before); + EXPECT_EQ(status.limit_bytes, expected_bytes - 1); + EXPECT_EQ(current_mega_moe_weight_cache_bytes(), before); +} + +TEST(MegaMoeRuntimeTest, ReturnsWeightCacheAccountingOnDestruction) { + constexpr int64_t kExperts = 2; + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + const auto canonical_w13 = + torch::empty({kExperts, 2 * kIntermediate, kHidden}, options); + const auto canonical_w2 = + torch::empty({kExperts, kHidden, kIntermediate}, options); + const int64_t expected_bytes = + (canonical_w13.numel() + canonical_w2.numel()) * + static_cast(canonical_w13.element_size()); + const int64_t before = current_mega_moe_weight_cache_bytes(); + + { + MegaMoeWeightBudgetStatus status; + auto cache = try_build_mega_moe_weight_cache(canonical_w13, + canonical_w2, + kHidden, + kIntermediate, + expected_bytes, + &status); + ASSERT_TRUE(cache.has_value()); + EXPECT_EQ(current_mega_moe_weight_cache_bytes(), + before + expected_bytes); + } + + EXPECT_EQ(current_mega_moe_weight_cache_bytes(), before); +} + +TEST(MegaMoeRuntimeTest, ReleasesTensorFieldsBeforeBudgetReservation) { + constexpr int64_t kExperts = 2; + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + const auto canonical_w13 = + torch::empty({kExperts, 2 * kIntermediate, kHidden}, options); + const auto canonical_w2 = + torch::empty({kExperts, kHidden, kIntermediate}, options); + const int64_t expected_bytes = + (canonical_w13.numel() + canonical_w2.numel()) * + static_cast(canonical_w13.element_size()); + std::vector events; + MegaMoeWeightCacheDestructionObserver observer; + observer.on_tensor_fields_released = + [&]() { events.push_back("tensor_fields_released"); }; + observer.on_reservation_released = + [&]() { events.push_back("reservation_released"); }; + + { + MegaMoeWeightBudgetStatus status; + auto cache = try_build_mega_moe_weight_cache(canonical_w13, + canonical_w2, + kHidden, + kIntermediate, + expected_bytes, + &status, + &observer); + ASSERT_TRUE(cache.has_value()); + EXPECT_TRUE(events.empty()); + } + + ASSERT_EQ(events.size(), 2); + EXPECT_EQ(events[0], "tensor_fields_released"); + EXPECT_EQ(events[1], "reservation_released"); +} + +TEST(MegaMoeRuntimeTest, RejectsCanonicalShapeMismatch) { + constexpr int64_t kHidden = 4; + constexpr int64_t kIntermediate = 2; + const auto options = torch::TensorOptions().dtype(torch::kBFloat16); + const torch::Tensor canonical_w13 = + torch::empty({3, 2 * kIntermediate, kHidden + 1}, options); + const torch::Tensor canonical_w2 = + torch::empty({3, kHidden, kIntermediate}, options); + + EXPECT_THROW(build_mega_moe_weight_cache( + canonical_w13, canonical_w2, kHidden, kIntermediate), + c10::Error); +} + +TEST(MegaMoeRuntimeTest, PreparesBfloat16TopkWeightsForA3MegaMoe) { + const torch::Tensor topk_weights = + torch::tensor({{0.125F, 0.25F, 0.625F}}, + torch::TensorOptions().dtype(torch::kBFloat16)); + + const torch::Tensor prepared = + prepare_mega_moe_topk_weights(topk_weights); + + EXPECT_EQ(prepared.scalar_type(), torch::kFloat32); + EXPECT_TRUE(prepared.is_contiguous()); + EXPECT_EQ(prepared.sizes(), topk_weights.sizes()); + EXPECT_TRUE(torch::equal(prepared, + topk_weights.to(torch::kFloat32))); + EXPECT_NE(prepared.data_ptr(), topk_weights.data_ptr()); +} + +} // namespace +} // namespace xllm::layer diff --git a/third_party/torch_npu_ops b/third_party/torch_npu_ops index adc29dc9d2..21cab2b6af 160000 --- a/third_party/torch_npu_ops +++ b/third_party/torch_npu_ops @@ -1 +1 @@ -Subproject commit adc29dc9d2e69875524b8a54bfc251463275d66e +Subproject commit 21cab2b6af14f18f56b7453024319065c8fb9764 diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index ecc6a13d03..498b3761e8 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -395,6 +395,10 @@ DECLARE_bool(enable_intralayer_addnorm); DECLARE_int32(enable_fused_mc2); +DECLARE_string(mega_moe_mode); + +DECLARE_int64(mega_moe_weight_cache_budget_bytes); + DECLARE_bool(enable_interlayer_addnorm); DECLARE_bool(enable_split_rmsnorm_rope); diff --git a/xllm/core/framework/config/kernel_config.cpp b/xllm/core/framework/config/kernel_config.cpp index bc6bf419fe..de0151e09b 100644 --- a/xllm/core/framework/config/kernel_config.cpp +++ b/xllm/core/framework/config/kernel_config.cpp @@ -38,6 +38,13 @@ DEFINE_int32(enable_fused_mc2, "MC2, positive values enable dense matmul-allreduce, 1 uses " "DispatchFFNCombine for MoE, 2 uses DispatchGmmCombineDecode for " "MoE."); +DEFINE_string(mega_moe_mode, + "off", + "MegaMoe mode for NPU: off, auto, or on. Default is off."); +DEFINE_int64(mega_moe_weight_cache_budget_bytes, + 4LL * 1024 * 1024 * 1024, + "Process-wide byte budget for independent contiguous MegaMoe " + "weight caches."); DEFINE_bool(enable_interlayer_addnorm, false, "enable fused interlayer addnorm ops."); @@ -71,6 +78,17 @@ int32_t resolve_fused_mc2_mode(int32_t mode) { } return 0; } + +void validate_mega_moe_mode(const std::string& mode) { + CHECK(mode == "off" || mode == "auto" || mode == "on") + << "--mega_moe_mode must be off, auto, or on; got " << mode; +} + +void validate_mega_moe_weight_cache_budget(int64_t budget_bytes) { + CHECK_GE(budget_bytes, 0) + << "--mega_moe_weight_cache_budget_bytes must be >= 0; got " + << budget_bytes; +} #endif } // namespace @@ -81,6 +99,8 @@ void KernelConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(npu_kernel_backend); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_intralayer_addnorm); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_fused_mc2); + XLLM_CONFIG_ASSIGN_FROM_FLAG(mega_moe_mode); + XLLM_CONFIG_ASSIGN_FROM_FLAG(mega_moe_weight_cache_budget_bytes); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_interlayer_addnorm); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_split_rmsnorm_rope); XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_aclnn_matmul); @@ -94,6 +114,8 @@ void KernelConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(npu_kernel_backend); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_intralayer_addnorm); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_fused_mc2); + XLLM_CONFIG_ASSIGN_FROM_JSON(mega_moe_mode); + XLLM_CONFIG_ASSIGN_FROM_JSON(mega_moe_weight_cache_budget_bytes); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_interlayer_addnorm); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_split_rmsnorm_rope); XLLM_CONFIG_ASSIGN_FROM_JSON(enable_aclnn_matmul); @@ -113,6 +135,10 @@ void KernelConfig::append_config_json( config_json, default_config, enable_intralayer_addnorm); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_fused_mc2); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, mega_moe_mode); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, mega_moe_weight_cache_budget_bytes); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, enable_interlayer_addnorm); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( @@ -136,6 +162,9 @@ void KernelConfig::initialize() { } #if defined(USE_NPU) enable_fused_mc2(resolve_fused_mc2_mode(enable_fused_mc2())); + validate_mega_moe_mode(mega_moe_mode()); + validate_mega_moe_weight_cache_budget( + mega_moe_weight_cache_budget_bytes()); #endif } diff --git a/xllm/core/framework/config/kernel_config.h b/xllm/core/framework/config/kernel_config.h index 6be36a3456..dbc9d34037 100644 --- a/xllm/core/framework/config/kernel_config.h +++ b/xllm/core/framework/config/kernel_config.h @@ -44,6 +44,8 @@ class KernelConfig final { "npu_kernel_backend", "enable_intralayer_addnorm", "enable_fused_mc2", + "mega_moe_mode", + "mega_moe_weight_cache_budget_bytes", "enable_interlayer_addnorm", "enable_split_rmsnorm_rope", "enable_aclnn_matmul", @@ -60,6 +62,11 @@ class KernelConfig final { PROPERTY(int32_t, enable_fused_mc2) = 0; + PROPERTY(std::string, mega_moe_mode) = "off"; + + PROPERTY(int64_t, mega_moe_weight_cache_budget_bytes) = + 4LL * 1024 * 1024 * 1024; + PROPERTY(bool, enable_interlayer_addnorm) = false; PROPERTY(bool, enable_split_rmsnorm_rope) = false; diff --git a/xllm/core/framework/parallel_state/CMakeLists.txt b/xllm/core/framework/parallel_state/CMakeLists.txt index 26bdab465d..9ae476e76a 100644 --- a/xllm/core/framework/parallel_state/CMakeLists.txt +++ b/xllm/core/framework/parallel_state/CMakeLists.txt @@ -1,5 +1,24 @@ include(cc_library) +if(USE_NPU) + cc_library( + NAME + mega_moe_comm_resource + HDRS + mega_moe_comm_resource.h + SRCS + mega_moe_comm_resource.cpp + DEPS + torch + torch_npu + hccl + hcomm + glog::glog + ) + target_link_libraries(mega_moe_comm_resource PRIVATE ${CMAKE_DL_LIBS}) + +endif() + cc_library( NAME parallel_state @@ -44,6 +63,7 @@ cc_library( $<$:hccl> $<$:hcomm> $<$:xllm_atb_layers> + $<$:mega_moe_comm_resource> gflags::gflags glog::glog ) diff --git a/xllm/core/framework/parallel_state/mega_moe_comm_resource.cpp b/xllm/core/framework/parallel_state/mega_moe_comm_resource.cpp new file mode 100644 index 0000000000..0d029f54cf --- /dev/null +++ b/xllm/core/framework/parallel_state/mega_moe_comm_resource.cpp @@ -0,0 +1,469 @@ +/* Copyright 2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "framework/parallel_state/mega_moe_comm_resource.h" + +#include +#include +#include + +#include +#include +#include +#include + +namespace xllm { +namespace { + +constexpr uint32_t kHcclMaxRankSize = 1024; +constexpr uint8_t kCommEngineAiv = 4; +constexpr uint8_t kAllToAllOpType = 8; +constexpr char kHcclLibrary[] = "libhccl.so"; +constexpr char kHcclFrameworkLibrary[] = "libhccl_fwk.so"; +constexpr char kAllToAllAlgorithm[] = + "AlltoAll=level0:fullmesh;level1:pairwise"; + +// This ABI is consumed by the CANN 9.1 MegaMoe kernel. It is intentionally +// kept local to the bridge so communication details do not leak into layers. +// The layout matches cann-ops-transformer commit +// 09f2ed7da10633ac704cf281d74849cc04cde7d9 comm_context.cpp (KFC mode). +struct MegaMoeCommContext { + uint32_t ep_rank_id = 0; + uint32_t rank_size_per_server = 0; + uint64_t kfc_context_addr = 0; + std::array ep_hccl_buffer = {}; + std::array hcomm_handle = {}; +}; + +static_assert(sizeof(MegaMoeCommContext) % sizeof(int32_t) == 0); + +using HcclKfcAllocOpArgs = HcclResult (*)(void**); +using HcclKfcFreeOpArgs = HcclResult (*)(void*); +using HcclKfcOpArgsSetAlgConfig = HcclResult (*)(void*, char*); +using HcclKfcOpArgsSetCommEngine = HcclResult (*)(void*, uint8_t); +using HcclCreateOpResCtx = HcclResult (*)(HcclComm, + uint8_t, + void*, + void**); +using HcclGetRankSize = HcclResult (*)(HcclComm, uint32_t*); +using HcclGetRankId = HcclResult (*)(HcclComm, uint32_t*); +using HcclGetHcclBuffer = HcclResult (*)(HcclComm, void**, uint64_t*); +using HcclGetRemoteIpcHcclBuf = HcclResult (*)(HcclComm, + uint64_t, + void**, + uint64_t*); + +template +Function load_symbol(void* library, const char* symbol) { + return reinterpret_cast(dlsym(library, symbol)); +} + +class MegaMoeCommApis final { + public: + MegaMoeCommApis() { + hccl_library_ = dlopen(kHcclLibrary, RTLD_NOW | RTLD_LOCAL); + if (hccl_library_ == nullptr) { + missing_symbol_ = kHcclLibrary; + return; + } + hccl_framework_library_ = + dlopen(kHcclFrameworkLibrary, RTLD_NOW | RTLD_LOCAL); + if (hccl_framework_library_ == nullptr) { + missing_symbol_ = kHcclFrameworkLibrary; + return; + } + + if (!load_hccl_symbol(kfc_alloc_op_args_, "HcclKfcAllocOpArgs") || + !load_hccl_symbol(kfc_free_op_args_, "HcclKfcFreeOpArgs") || + !load_hccl_symbol(kfc_set_alg_config_, + "HcclKfcOpArgsSetAlgConfig") || + !load_hccl_symbol(kfc_set_comm_engine_, + "HcclKfcOpArgsSetCommEngine") || + !load_hccl_symbol(create_op_res_ctx_, "HcclCreateOpResCtx") || + !load_hccl_symbol(get_rank_size_, "HcclGetRankSize") || + !load_hccl_symbol(get_rank_id_, "HcclGetRankId") || + !load_hccl_symbol(get_hccl_buffer_, "HcclGetHcclBuffer") || + !load_framework_symbol(get_remote_ipc_hccl_buffer_, + "HcclGetRemoteIpcHcclBuf")) { + return; + } + available_ = true; + } + + MegaMoeCommApis(const MegaMoeCommApis&) = delete; + MegaMoeCommApis& operator=(const MegaMoeCommApis&) = delete; + + bool available() const { return available_; } + const std::string& missing_symbol() const { return missing_symbol_; } + + HcclKfcAllocOpArgs kfc_alloc_op_args() const { + return kfc_alloc_op_args_; + } + HcclKfcFreeOpArgs kfc_free_op_args() const { + return kfc_free_op_args_; + } + HcclKfcOpArgsSetAlgConfig kfc_set_alg_config() const { + return kfc_set_alg_config_; + } + HcclKfcOpArgsSetCommEngine kfc_set_comm_engine() const { + return kfc_set_comm_engine_; + } + HcclCreateOpResCtx create_op_res_ctx() const { + return create_op_res_ctx_; + } + HcclGetRankSize get_rank_size() const { return get_rank_size_; } + HcclGetRankId get_rank_id() const { return get_rank_id_; } + HcclGetHcclBuffer get_hccl_buffer() const { return get_hccl_buffer_; } + HcclGetRemoteIpcHcclBuf get_remote_ipc_hccl_buffer() const { + return get_remote_ipc_hccl_buffer_; + } + + private: + template + bool load_hccl_symbol(Function& function, const char* symbol) { + function = load_symbol(hccl_library_, symbol); + return record_missing_symbol(function, symbol); + } + + template + bool load_framework_symbol(Function& function, const char* symbol) { + function = load_symbol(hccl_framework_library_, symbol); + return record_missing_symbol(function, symbol); + } + + template + bool record_missing_symbol(Function function, const char* symbol) { + if (function != nullptr) { + return true; + } + missing_symbol_ = symbol; + return false; + } + + // CANN communication libraries are process-wide runtimes. Keep their + // handles loaded until process exit; late dlclose can race HCCL teardown. + void* hccl_library_ = nullptr; + void* hccl_framework_library_ = nullptr; + bool available_ = false; + std::string missing_symbol_; + HcclKfcAllocOpArgs kfc_alloc_op_args_ = nullptr; + HcclKfcFreeOpArgs kfc_free_op_args_ = nullptr; + HcclKfcOpArgsSetAlgConfig kfc_set_alg_config_ = nullptr; + HcclKfcOpArgsSetCommEngine kfc_set_comm_engine_ = nullptr; + HcclCreateOpResCtx create_op_res_ctx_ = nullptr; + HcclGetRankSize get_rank_size_ = nullptr; + HcclGetRankId get_rank_id_ = nullptr; + HcclGetHcclBuffer get_hccl_buffer_ = nullptr; + HcclGetRemoteIpcHcclBuf get_remote_ipc_hccl_buffer_ = nullptr; +}; + +MegaMoeCommApis& comm_apis() { + static MegaMoeCommApis apis; + return apis; +} + +class OpArgsGuard final { + public: + OpArgsGuard(void* op_args, HcclKfcFreeOpArgs free_op_args) + : op_args_(op_args), free_op_args_(free_op_args) {} + + ~OpArgsGuard() { + if (op_args_ == nullptr) { + return; + } + const HcclResult result = free_op_args_(op_args_); + if (result != HCCL_SUCCESS) { + LOG(ERROR) << "HcclKfcFreeOpArgs failed during cleanup, result=" + << static_cast(result); + } + } + + OpArgsGuard(const OpArgsGuard&) = delete; + OpArgsGuard& operator=(const OpArgsGuard&) = delete; + + void release() { op_args_ = nullptr; } + + private: + void* op_args_ = nullptr; + HcclKfcFreeOpArgs free_op_args_ = nullptr; +}; + +void check_hccl_result(HcclResult result, const char* operation) { + CHECK_EQ(result, HCCL_SUCCESS) + << operation << " failed, result=" << static_cast(result); +} + +MegaMoeCommContext build_context(const MegaMoeCommSpec& spec, + int64_t& ccl_buffer_size) { + MegaMoeCommApis& apis = comm_apis(); + CHECK(apis.available()) << "missing MegaMoe communication symbol: " + << apis.missing_symbol(); + + void* op_args = nullptr; + check_hccl_result(apis.kfc_alloc_op_args()(&op_args), + "HcclKfcAllocOpArgs"); + CHECK(op_args != nullptr) << "HcclKfcAllocOpArgs returned null"; + OpArgsGuard op_args_guard(op_args, apis.kfc_free_op_args()); + + check_hccl_result( + apis.kfc_set_comm_engine()(op_args, kCommEngineAiv), + "HcclKfcOpArgsSetCommEngine"); + check_hccl_result(apis.kfc_set_alg_config()( + op_args, const_cast(kAllToAllAlgorithm)), + "HcclKfcOpArgsSetAlgConfig"); + + MegaMoeCommContext context; + void* op_resource_context = nullptr; + check_hccl_result(apis.create_op_res_ctx()(spec.hccl_comm, + kAllToAllOpType, + op_args, + &op_resource_context), + "HcclCreateOpResCtx"); + CHECK(op_resource_context != nullptr) << "HcclCreateOpResCtx returned null"; + context.kfc_context_addr = + reinterpret_cast(op_resource_context); + + const HcclResult free_result = apis.kfc_free_op_args()(op_args); + op_args_guard.release(); + check_hccl_result(free_result, "HcclKfcFreeOpArgs"); + + uint32_t rank_size = 0; + uint32_t rank_id = 0; + check_hccl_result(apis.get_rank_size()(spec.hccl_comm, &rank_size), + "HcclGetRankSize"); + check_hccl_result(apis.get_rank_id()(spec.hccl_comm, &rank_id), + "HcclGetRankId"); + CHECK_EQ(rank_size, static_cast(spec.ep_world_size)); + CHECK_LT(rank_id, rank_size); + context.ep_rank_id = rank_id; + + std::vector rank_buffer_addresses(rank_size, nullptr); + std::vector rank_accessible_spans(rank_size, 0); + for (uint32_t remote_rank_id = 0; remote_rank_id < rank_size; + ++remote_rank_id) { + void* remote_address = nullptr; + uint64_t remote_size = 0; + const char* buffer_query = nullptr; + HcclResult result = HCCL_E_INTERNAL; + if (remote_rank_id == rank_id) { + result = apis.get_hccl_buffer()( + spec.hccl_comm, &remote_address, &remote_size); + ccl_buffer_size = static_cast(remote_size); + buffer_query = "HcclGetHcclBuffer"; + } else { + result = apis.get_remote_ipc_hccl_buffer()(spec.hccl_comm, + remote_rank_id, + &remote_address, + &remote_size); + buffer_query = "HcclGetRemoteIpcHcclBuf"; + } + check_hccl_result(result, buffer_query); + CHECK(remote_address != nullptr); + rank_buffer_addresses[remote_rank_id] = remote_address; + rank_accessible_spans[remote_rank_id] = remote_size; + } + + const auto span_validation = validate_mega_moe_buffer_accessible_spans( + static_cast(ccl_buffer_size), rank_accessible_spans); + CHECK(span_validation.valid) + << "MegaMoe HCCL accessible span is smaller than the local payload at " + "rank " + << span_validation.mismatched_rank + << ": required_payload=" << span_validation.required_payload_size + << ", accessible_span=" << span_validation.accessible_span; + + for (uint32_t remote_rank_id = 0; remote_rank_id < rank_size; + ++remote_rank_id) { + context.ep_hccl_buffer[remote_rank_id] = + reinterpret_cast(rank_buffer_addresses[remote_rank_id]); + } + return context; +} + +} // namespace + +MegaMoeCommValidation validate_mega_moe_comm_spec( + const MegaMoeCommSpec& spec) { + if (spec.group_name.empty()) { + return {false, MegaMoeCommRejectReason::EMPTY_GROUP}; + } + if (spec.hccl_comm == nullptr) { + return {false, MegaMoeCommRejectReason::NULL_COMM}; + } + if (spec.ep_world_size <= 0 || + static_cast(spec.ep_world_size) > kHcclMaxRankSize) { + return {false, MegaMoeCommRejectReason::UNSUPPORTED_EP_WORLD_SIZE}; + } + if (spec.device_index < 0) { + return {false, MegaMoeCommRejectReason::INVALID_DEVICE_INDEX}; + } + if (spec.max_num_tokens_per_rank <= 0) { + return {false, + MegaMoeCommRejectReason::INVALID_MAX_NUM_TOKENS_PER_RANK}; + } + return {true, MegaMoeCommRejectReason::NONE}; +} + +MegaMoeCommSymbolStatus probe_mega_moe_comm_symbols() { + MegaMoeCommApis& apis = comm_apis(); + return {apis.available(), apis.missing_symbol()}; +} + +MegaMoeBufferSpanValidation validate_mega_moe_buffer_accessible_spans( + uint64_t local_payload_size, + const std::vector& rank_accessible_spans) { + if (local_payload_size == 0) { + return {false, -1, local_payload_size, 0}; + } + if (rank_accessible_spans.empty()) { + return {false, -1, local_payload_size, 0}; + } + for (size_t rank = 0; rank < rank_accessible_spans.size(); ++rank) { + const uint64_t accessible_span = rank_accessible_spans[rank]; + if (accessible_span < local_payload_size) { + return {false, + static_cast(rank), + local_payload_size, + accessible_span}; + } + } + return {true, -1, local_payload_size, 0}; +} + +class MegaMoeCommResource::Impl final { + public: + explicit Impl(const MegaMoeCommSpec& spec) + : max_num_tokens_per_rank_(spec.max_num_tokens_per_rank) { + const MegaMoeCommContext context = build_context(spec, ccl_buffer_size_); + const int64_t context_elements = + static_cast(sizeof(context) / sizeof(int32_t)); + host_context_tensor_ = at::empty({context_elements}, + at::TensorOptions().dtype(at::kInt)); + std::memcpy(host_context_tensor_.data_ptr(), + &context, + sizeof(context)); + context_tensor_ = host_context_tensor_.to( + at::Device(at::kPrivateUse1, spec.device_index), + at::kInt, + false, + true); + } + + const at::Tensor& context_tensor() const { return context_tensor_; } + int64_t ccl_buffer_size() const { return ccl_buffer_size_; } + int64_t max_num_tokens_per_rank() const { + return max_num_tokens_per_rank_; + } + + private: + // op_resource_context is owned by the non-owning HCCL communicator. Keep + // only tensor data and destroy it before ParallelArgs tears that comm down. + at::Tensor host_context_tensor_; + at::Tensor context_tensor_; + int64_t ccl_buffer_size_ = 0; + int64_t max_num_tokens_per_rank_ = 0; +}; + +MegaMoeCommResource::MegaMoeCommResource(std::unique_ptr impl) + : impl_(std::move(impl)) {} + +MegaMoeCommResource::~MegaMoeCommResource() = default; + +std::unique_ptr MegaMoeCommResource::create( + const MegaMoeCommSpec& spec) { + const MegaMoeCommValidation validation = validate_mega_moe_comm_spec(spec); + CHECK(validation.valid) + << "invalid MegaMoe communication spec, reason=" + << static_cast(validation.reason); + CHECK(probe_mega_moe_comm_symbols().available) + << "MegaMoe communication symbols are unavailable"; + return std::unique_ptr( + new MegaMoeCommResource(std::make_unique(spec))); +} + +const at::Tensor& MegaMoeCommResource::context_tensor() const { + return impl_->context_tensor(); +} + +int64_t MegaMoeCommResource::ccl_buffer_size() const { + return impl_->ccl_buffer_size(); +} + +int64_t MegaMoeCommResource::max_num_tokens_per_rank() const { + return impl_->max_num_tokens_per_rank(); +} + +bool MegaMoeCommResourceSlot::same_key(const MegaMoeCommSpec& lhs, + const MegaMoeCommSpec& rhs) { + return lhs.hccl_comm == rhs.hccl_comm && + lhs.group_name == rhs.group_name && + lhs.ep_world_size == rhs.ep_world_size && + lhs.device_index == rhs.device_index && + lhs.max_num_tokens_per_rank == rhs.max_num_tokens_per_rank; +} + +std::shared_ptr MegaMoeCommResourceSlot::acquire( + const MegaMoeCommSpec& spec) { + return acquire(spec, [](const MegaMoeCommSpec& create_spec) { + return std::shared_ptr( + MegaMoeCommResource::create(create_spec)); + }); +} + +std::shared_ptr MegaMoeCommResourceSlot::acquire( + const MegaMoeCommSpec& spec, + const Factory& factory) { + std::lock_guard lock(mutex_); + if (resource_ != nullptr && cached_spec_.has_value() && + same_key(cached_spec_.value(), spec)) { + return resource_; + } + + // Run the factory before mutating the slot. A throw therefore leaves the + // previous entry intact and a first-create failure cannot be cached. + std::shared_ptr candidate = factory(spec); + CHECK(candidate != nullptr) + << "MegaMoe communication resource factory returned null."; + cached_spec_ = spec; + resource_ = std::move(candidate); + return resource_; +} + +void MegaMoeCommResourceSlot::reset() { + // Destroy outside the lock. The resource owns an NPU tensor and its + // destructor may synchronize runtime state. + std::shared_ptr released; + { + std::lock_guard lock(mutex_); + cached_spec_.reset(); + released = std::move(resource_); + } +} + +void MegaMoeCommResourceSlot::reset_for_teardown() { + // A layer owns only a weak_ptr. Any additional strong owner is therefore a + // forward that has locked the resource for the duration of its collective. + std::shared_ptr released; + { + std::lock_guard lock(mutex_); + CHECK(resource_ == nullptr || resource_.use_count() == 1) + << "cannot tear down MegaMoe communication resource while an " + "active MegaMoe forward still holds it"; + cached_spec_.reset(); + released = std::move(resource_); + } +} + +} // namespace xllm diff --git a/xllm/core/framework/parallel_state/mega_moe_comm_resource.h b/xllm/core/framework/parallel_state/mega_moe_comm_resource.h new file mode 100644 index 0000000000..d8cb4c9d36 --- /dev/null +++ b/xllm/core/framework/parallel_state/mega_moe_comm_resource.h @@ -0,0 +1,129 @@ +/* Copyright 2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include +#include +#include +#include +#include + +#include "hccl/hccl.h" + +namespace at { +class Tensor; +} // namespace at + +namespace xllm { + +struct MegaMoeCommSpec { + std::string group_name; + HcclComm hccl_comm = nullptr; + int32_t ep_world_size = 0; + int32_t device_index = -1; + int64_t max_num_tokens_per_rank = 0; +}; + +enum class MegaMoeCommRejectReason : int32_t { + NONE = 0, + EMPTY_GROUP, + NULL_COMM, + UNSUPPORTED_EP_WORLD_SIZE, + INVALID_DEVICE_INDEX, + INVALID_MAX_NUM_TOKENS_PER_RANK, +}; + +struct MegaMoeCommValidation { + bool valid = false; + MegaMoeCommRejectReason reason = MegaMoeCommRejectReason::EMPTY_GROUP; +}; + +struct MegaMoeCommSymbolStatus { + bool available = false; + std::string missing_symbol; +}; + +MegaMoeCommValidation validate_mega_moe_comm_spec( + const MegaMoeCommSpec& spec); + +MegaMoeCommSymbolStatus probe_mega_moe_comm_symbols(); + +struct MegaMoeBufferSpanValidation { + bool valid = false; + int32_t mismatched_rank = -1; + uint64_t required_payload_size = 0; + uint64_t accessible_span = 0; +}; + +// HcclGetHcclBuffer returns the local communication payload while +// HcclGetRemoteIpcHcclBuf may return a larger accessible span that includes +// auxiliary MC2 memory. Validate that every returned span covers the complete +// local payload before copying the communication context to an NPU tensor. +MegaMoeBufferSpanValidation validate_mega_moe_buffer_accessible_spans( + uint64_t local_payload_size, + const std::vector& rank_accessible_spans); + +class MegaMoeCommResource final { + public: + ~MegaMoeCommResource(); + + MegaMoeCommResource(const MegaMoeCommResource&) = delete; + MegaMoeCommResource& operator=(const MegaMoeCommResource&) = delete; + MegaMoeCommResource(MegaMoeCommResource&&) = delete; + MegaMoeCommResource& operator=(MegaMoeCommResource&&) = delete; + + static std::unique_ptr create( + const MegaMoeCommSpec& spec); + + const at::Tensor& context_tensor() const; + int64_t ccl_buffer_size() const; + int64_t max_num_tokens_per_rank() const; + private: + class Impl; + + explicit MegaMoeCommResource(std::unique_ptr impl); + + std::unique_ptr impl_; +}; + +// A bounded, ProcessGroup-owned cache. One strong entry follows the HCCL +// communicator lifetime; layers retain only weak references. There is no +// process-global unbounded key map. +class MegaMoeCommResourceSlot final { + public: + using Factory = std::function( + const MegaMoeCommSpec&)>; + + std::shared_ptr acquire( + const MegaMoeCommSpec& spec); + std::shared_ptr acquire( + const MegaMoeCommSpec& spec, + const Factory& factory); + void reset(); + void reset_for_teardown(); + + private: + static bool same_key(const MegaMoeCommSpec& lhs, + const MegaMoeCommSpec& rhs); + + std::mutex mutex_; + std::optional cached_spec_; + std::shared_ptr resource_; +}; + +} // namespace xllm diff --git a/xllm/core/framework/parallel_state/npu_process_group.cpp b/xllm/core/framework/parallel_state/npu_process_group.cpp index 371dc472d5..d38bc55226 100644 --- a/xllm/core/framework/parallel_state/npu_process_group.cpp +++ b/xllm/core/framework/parallel_state/npu_process_group.cpp @@ -169,6 +169,7 @@ ProcessGroupImpl::ProcessGroupImpl(int32_t global_rank, // Destructor. ProcessGroupImpl::~ProcessGroupImpl() { + release_mega_moe_comm_resource(); if (pg_) { shutdown_backend(); } else if (comm_ != nullptr) { @@ -199,4 +200,22 @@ std::string ProcessGroupImpl::hccl_comm_name(bool init_comm) { #endif } +HcclComm ProcessGroupImpl::hccl_comm() { + if (comm_ != nullptr) { + return comm_; + } + CHECK(pg_ != nullptr) << "HCCL communicator requires an initialized group."; +#if defined(USE_NPU) && \ + (TORCH_VERSION_MAJOR < 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR < 7)) + const int64_t comm_handle = pg_->getHcclComm(pg_->getRank()); +#else + auto* hccl_pg = dynamic_cast(pg_.get()); + CHECK(hccl_pg != nullptr) << "Process group is not NPU HCCL."; + const int64_t comm_handle = hccl_pg->getHcclComm(pg_->getRank()); +#endif + CHECK(comm_handle != 0) << "NPU HCCL process group returned a null comm."; + return reinterpret_cast(comm_handle); +} + } // namespace xllm diff --git a/xllm/core/framework/parallel_state/npu_process_group.h b/xllm/core/framework/parallel_state/npu_process_group.h index 28aaf3aa63..47896d78af 100644 --- a/xllm/core/framework/parallel_state/npu_process_group.h +++ b/xllm/core/framework/parallel_state/npu_process_group.h @@ -57,6 +57,8 @@ class ProcessGroupImpl : public ProcessGroup { std::string hccl_comm_name(bool init_comm = true) override; + HcclComm hccl_comm() override; + private: HcclComm comm_ = nullptr; c10_npu::NPUStream comm_stream_; diff --git a/xllm/core/framework/parallel_state/process_group.cpp b/xllm/core/framework/parallel_state/process_group.cpp index 73f91ecb4d..8687930149 100644 --- a/xllm/core/framework/parallel_state/process_group.cpp +++ b/xllm/core/framework/parallel_state/process_group.cpp @@ -16,6 +16,7 @@ limitations under the License. #include "process_group.h" #if defined(USE_NPU) +#include "mega_moe_comm_resource.h" #include "npu_process_group.h" #elif defined(USE_MLU) #include "mlu_process_group.h" @@ -58,7 +59,22 @@ std::vector get_gather_shape(int32_t world_size, namespace xllm { -ProcessGroup::~ProcessGroup() { shutdown_backend(); } +ProcessGroup::ProcessGroup(int32_t rank, + int32_t world_size, + const torch::Device& device) + : rank_(rank), world_size_(world_size), device_(device) { +#if defined(USE_NPU) + mega_moe_comm_resource_slot_ = + std::make_unique(); +#endif +} + +ProcessGroup::~ProcessGroup() { +#if defined(USE_NPU) + release_mega_moe_comm_resource(); +#endif + shutdown_backend(); +} void ProcessGroup::shutdown_backend() { if (pg_ == nullptr) { @@ -68,6 +84,23 @@ void ProcessGroup::shutdown_backend() { pg_.reset(); } +#if defined(USE_NPU) +std::shared_ptr +ProcessGroup::acquire_mega_moe_comm_resource( + const MegaMoeCommSpec& spec) { + CHECK(mega_moe_comm_resource_slot_ != nullptr) + << "MegaMoe communication resource slot is unavailable."; + return mega_moe_comm_resource_slot_->acquire(spec); +} + +void ProcessGroup::release_mega_moe_comm_resource() { + if (mega_moe_comm_resource_slot_ != nullptr) { + mega_moe_comm_resource_slot_->reset_for_teardown(); + mega_moe_comm_resource_slot_.reset(); + } +} +#endif + std::pair> get_group_rank(int world_size, int global_rank, int split_size, @@ -229,6 +262,13 @@ std::string ProcessGroup::hccl_comm_name(bool init_comm) { return ""; } +#if defined(USE_NPU) +HcclComm ProcessGroup::hccl_comm() { + CHECK(false) << "hccl_comm is only supported on NPU HCCL process group."; + return nullptr; +} +#endif + std::unique_ptr create_process_group( int32_t rank, int32_t world_size, diff --git a/xllm/core/framework/parallel_state/process_group.h b/xllm/core/framework/parallel_state/process_group.h index 5ed19fe9c6..ca8c1d5b57 100644 --- a/xllm/core/framework/parallel_state/process_group.h +++ b/xllm/core/framework/parallel_state/process_group.h @@ -29,6 +29,11 @@ limitations under the License. namespace xllm { class ProcessGroupImpl; +#if defined(USE_NPU) +struct MegaMoeCommSpec; +class MegaMoeCommResource; +class MegaMoeCommResourceSlot; +#endif std::pair> get_group_rank(int world_size, int global_rank, @@ -41,8 +46,7 @@ c10::intrusive_ptr create_tcp_store(const std::string& host, class ProcessGroup { public: - ProcessGroup(int32_t rank, int32_t world_size, const torch::Device& device) - : rank_(rank), world_size_(world_size), device_(device) {} + ProcessGroup(int32_t rank, int32_t world_size, const torch::Device& device); virtual ~ProcessGroup(); @@ -115,6 +119,12 @@ class ProcessGroup { virtual std::string hccl_comm_name(bool init_comm = true); +#if defined(USE_NPU) + virtual HcclComm hccl_comm(); + std::shared_ptr acquire_mega_moe_comm_resource( + const MegaMoeCommSpec& spec); +#endif + private: // rank of current process. int32_t rank_ = 0; @@ -125,8 +135,16 @@ class ProcessGroup { // device of current process. torch::Device device_; +#if defined(USE_NPU) + // Bounded to one cache entry and explicitly released before HCCL teardown. + std::unique_ptr mega_moe_comm_resource_slot_; +#endif + protected: void shutdown_backend(); +#if defined(USE_NPU) + void release_mega_moe_comm_resource(); +#endif #if defined(USE_NPU) && \ (TORCH_VERSION_MAJOR < 2 || \ diff --git a/xllm/core/kernels/npu/CMakeLists.txt b/xllm/core/kernels/npu/CMakeLists.txt index c647a19819..43d70e0ec3 100644 --- a/xllm/core/kernels/npu/CMakeLists.txt +++ b/xllm/core/kernels/npu/CMakeLists.txt @@ -23,6 +23,7 @@ cc_library( npu_grouped_matmul.cpp npu_moe_gating_topk_softmax.cpp npu_dispatch_ffn_combine.cpp + npu_mega_moe.cpp npu_dispatch_gmm_combine_decode.cpp npu_moe_distribute_combine_v2.cpp npu_moe_distribute_dispatch_v2.cpp diff --git a/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp b/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp index ce23958d54..48cb0cee65 100644 --- a/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp +++ b/xllm/core/kernels/npu/aclnn/pytorch_npu_helper.hpp @@ -522,6 +522,14 @@ inline aclTensorList* convert_type(const at::TensorList& at_tensor_list) { return acl_tensor_list; } +inline aclTensorList* convert_type( + const c10::optional& opt_tensor_list) { + if (opt_tensor_list.has_value()) { + return convert_type(opt_tensor_list.value()); + } + return nullptr; +} + inline aclTensor* convert_type(const c10::optional& opt_tensor) { if (opt_tensor.has_value() && opt_tensor.value().defined()) { return convert_type(opt_tensor.value()); @@ -609,6 +617,9 @@ inline void release(aclBoolArray* p) { } inline void release(aclTensorList* p) { + if (p == nullptr) { + return; + } static const auto acl_destroy_tensor_list = get_op_api_func("aclDestroyTensorList"); if (acl_destroy_tensor_list == nullptr) { diff --git a/xllm/core/kernels/npu/mega_moe_acl_contract.h b/xllm/core/kernels/npu/mega_moe_acl_contract.h new file mode 100644 index 0000000000..ae66713ac8 --- /dev/null +++ b/xllm/core/kernels/npu/mega_moe_acl_contract.h @@ -0,0 +1,82 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +#pragma once + +#include + +#include +#include + +namespace xllm::kernel::npu { + +// Deployed custom-vendor ABI from ops-transformer commit 09f2ed7. +template +void execute_mega_moe_acl_contract( + Execute&& execute, + torch::Tensor& context, + torch::Tensor& x, + torch::Tensor& topk_ids, + torch::Tensor& topk_weights, + torch::TensorList& weight1, + torch::TensorList& weight2, + std::optional& weight_scales1, + std::optional& weight_scales2, + std::optional& bias1, + std::optional& bias2, + std::optional& x_active_mask, + int64_t& moe_expert_num, + int64_t& ep_world_size, + int64_t& ccl_buffer_size, + int64_t& max_recv_token_num, + int64_t& dispatch_quant_mode, + int64_t& dispatch_quant_out_dtype, + int64_t& combine_quant_mode, + char*& comm_alg, + int64_t& num_max_tokens_per_rank, + char*& activation, + float& activation_clamp, + int64_t& topo_type, + int64_t& rank_num_per_server, + torch::Tensor& output, + torch::Tensor& expert_token_nums) { + std::forward(execute)(context, + x, + topk_ids, + topk_weights, + weight1, + weight2, + weight_scales1, + weight_scales2, + bias1, + bias2, + x_active_mask, + moe_expert_num, + ep_world_size, + ccl_buffer_size, + max_recv_token_num, + dispatch_quant_mode, + dispatch_quant_out_dtype, + combine_quant_mode, + comm_alg, + num_max_tokens_per_rank, + activation, + activation_clamp, + topo_type, + rank_num_per_server, + output, + expert_token_nums); +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_mega_moe.cpp b/xllm/core/kernels/npu/npu_mega_moe.cpp new file mode 100644 index 0000000000..3e99ded53d --- /dev/null +++ b/xllm/core/kernels/npu/npu_mega_moe.cpp @@ -0,0 +1,302 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include + +#include "core/kernels/npu/mega_moe_acl_contract.h" +#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "core/kernels/npu/npu_ops_api.h" + +namespace xllm::kernel::npu { +namespace { + +std::string symbol_library_path(void* symbol) { + if (symbol == nullptr) { + return {}; + } + Dl_info info{}; + if (dladdr(symbol, &info) == 0 || info.dli_fname == nullptr) { + return {}; + } + return aclnn::detail::real_path(info.dli_fname); +} + +std::string expected_custom_mega_moe_library() { + const char* custom_opp_path = std::getenv("ASCEND_CUSTOM_OPP_PATH"); + if (custom_opp_path == nullptr) { + return {}; + } + + std::string root(custom_opp_path); + // MegaMoe's ABI differs from the stock CANN 9.1 op. The validated MegaMoe + // vendor must be first, while standard companion vendors may follow for + // other Qwen3.5 ops. The resolved MegaMoe symbols are still required to + // come from this first vendor's single libcust_opapi.so. + const size_t separator = root.find(':'); + if (separator != std::string::npos) { + root.resize(separator); + } + if (root.empty()) { + return {}; + } + return aclnn::detail::real_path( + root + "/op_api/lib/" + aclnn::detail::get_cust_op_api_lib_name()); +} + +void validate_weight_lists(const torch::Tensor& x, + const torch::TensorList weight1, + const torch::TensorList weight2, + int64_t local_expert_num) { + TORCH_CHECK(!weight1.empty(), "MegaMoe expects non-empty weight1."); + TORCH_CHECK(!weight2.empty(), "MegaMoe expects non-empty weight2."); + TORCH_CHECK(weight1.size() == weight2.size(), + "MegaMoe weight1/weight2 list size mismatch: ", + weight1.size(), + " vs ", + weight2.size()); + TORCH_CHECK(static_cast(weight1.size()) == local_expert_num, + "MegaMoe expects one weight pair per local expert: ", + local_expert_num, + ", got ", + weight1.size()); + + const int64_t hidden_size = x.size(1); + for (size_t expert = 0; expert < weight1.size(); ++expert) { + const auto& w1 = weight1[expert]; + const auto& w2 = weight2[expert]; + TORCH_CHECK(w1.dim() == 2 && w2.dim() == 2, + "MegaMoe expert weights must be 2D at local expert ", + expert, + "."); + TORCH_CHECK(w1.scalar_type() == at::kBFloat16 && + w2.scalar_type() == at::kBFloat16, + "MegaMoe A16W16 expects bf16 weights at local expert ", + expert, + "."); + TORCH_CHECK(w1.size(0) == hidden_size, + "MegaMoe W1 hidden dimension mismatch at local expert ", + expert, + ": expected ", + hidden_size, + ", got ", + w1.size(0)); + TORCH_CHECK(w2.size(1) == hidden_size, + "MegaMoe W2 hidden dimension mismatch at local expert ", + expert, + ": expected ", + hidden_size, + ", got ", + w2.size(1)); + TORCH_CHECK(w1.size(1) == 2 * w2.size(0), + "MegaMoe W1/W2 intermediate dimension mismatch at local " + "expert ", + expert, + ": W1 columns ", + w1.size(1), + ", W2 rows ", + w2.size(0)); + } +} + +} // namespace + +MegaMoeOpApiProvenance validate_mega_moe_op_api_paths( + const std::string& expected_library, + const std::string& workspace_library, + const std::string& execute_library) { + MegaMoeOpApiProvenance provenance; + provenance.expected_library = expected_library; + provenance.workspace_library = workspace_library; + provenance.execute_library = execute_library; + provenance.same_library = !workspace_library.empty() && + workspace_library == execute_library; + provenance.compatible = !expected_library.empty() && + provenance.same_library && + workspace_library == expected_library; + return provenance; +} + +MegaMoeOpApiProvenance inspect_mega_moe_op_api_provenance() { + void* workspace = + aclnn::detail::get_op_api_func_addr("aclnnMegaMoeGetWorkspaceSize"); + void* execute = aclnn::detail::get_op_api_func_addr("aclnnMegaMoe"); + return validate_mega_moe_op_api_paths(expected_custom_mega_moe_library(), + symbol_library_path(workspace), + symbol_library_path(execute)); +} + +bool has_mega_moe() { + static const MegaMoeOpApiProvenance provenance = + inspect_mega_moe_op_api_provenance(); + return provenance.compatible; +} + +std::tuple apply_npu_mega_moe( + const torch::Tensor& context, + const torch::Tensor& x, + const torch::Tensor& topk_ids, + const torch::Tensor& topk_weights, + const torch::TensorList weight1, + const torch::TensorList weight2, + int64_t moe_expert_num, + int64_t ep_world_size, + int64_t ccl_buffer_size, + const std::optional& weight_scales1, + const std::optional& weight_scales2, + const std::optional& bias1, + const std::optional& bias2, + const std::optional& x_active_mask, + int64_t max_recv_token_num, + int64_t dispatch_quant_mode, + int64_t combine_quant_mode, + const std::string& comm_alg, + int64_t num_max_tokens_per_rank, + const std::string& activation, + float activation_clamp, + int64_t dispatch_quant_out_dtype, + int64_t topo_type, + int64_t rank_num_per_server) { + const auto provenance = inspect_mega_moe_op_api_provenance(); + TORCH_CHECK( + provenance.compatible, + "aclnnMegaMoe custom ABI provenance check failed. Expected both " + "symbols in '", + provenance.expected_library, + "', but GetWorkspaceSize resolved to '", + provenance.workspace_library, + "' and execute resolved to '", + provenance.execute_library, + "'. Source CANN 9.1 and put the validated MegaMoe vendor first in " + "ASCEND_CUSTOM_OPP_PATH; standard companion vendors may follow, but " + "both MegaMoe symbols must resolve to the first vendor."); + TORCH_CHECK(context.defined(), "MegaMoe expects a defined context tensor."); + TORCH_CHECK(context.dim() == 1, "MegaMoe expects 1D context."); + TORCH_CHECK(context.scalar_type() == at::kInt, + "MegaMoe expects int32 context, got ", + c10::toString(context.scalar_type())); + TORCH_CHECK(x.dim() == 2, "MegaMoe expects 2D x."); + TORCH_CHECK(x.scalar_type() == at::kBFloat16, + "MegaMoe verified A16W16 path expects bf16 x, got ", + c10::toString(x.scalar_type())); + TORCH_CHECK(topk_ids.dim() == 2, "MegaMoe expects 2D topk_ids."); + TORCH_CHECK(topk_ids.scalar_type() == at::kInt, + "MegaMoe expects int32 topk_ids, got ", + c10::toString(topk_ids.scalar_type())); + TORCH_CHECK(topk_weights.dim() == 2, + "MegaMoe expects 2D topk_weights."); + TORCH_CHECK(topk_weights.scalar_type() == at::kFloat, + "MegaMoe verified A3 path requires fp32 topk_weights, got ", + c10::toString(topk_weights.scalar_type())); + TORCH_CHECK(topk_ids.sizes() == topk_weights.sizes(), + "MegaMoe topk_ids/topk_weights shape mismatch: ", + topk_ids.sizes(), + " vs ", + topk_weights.sizes()); + TORCH_CHECK(topk_ids.size(0) == x.size(0), + "MegaMoe x/router token count mismatch: ", + x.size(0), + " vs ", + topk_ids.size(0)); + TORCH_CHECK(moe_expert_num > 0, + "MegaMoe requires moe_expert_num > 0."); + TORCH_CHECK(ep_world_size > 0, "MegaMoe requires ep_world_size > 0."); + TORCH_CHECK(moe_expert_num % ep_world_size == 0, + "MegaMoe moe_expert_num must be divisible by ep_world_size: ", + moe_expert_num, + " vs ", + ep_world_size); + TORCH_CHECK(ccl_buffer_size > 0, + "MegaMoe requires ccl_buffer_size > 0."); + + const int64_t local_expert_num = moe_expert_num / ep_world_size; + validate_weight_lists(x, weight1, weight2, local_expert_num); + + TORCH_CHECK(!weight_scales1.has_value() && !weight_scales2.has_value() && + !bias1.has_value() && !bias2.has_value() && + !x_active_mask.has_value(), + "MegaMoe verified A16W16 path does not accept scales, biases, " + "or x_active_mask."); + TORCH_CHECK(dispatch_quant_mode == 0 && combine_quant_mode == 0, + "MegaMoe verified A16W16 path requires quant modes 0."); + TORCH_CHECK(activation == "swiglu", + "MegaMoe verified path requires swiglu activation, got ", + activation); + TORCH_CHECK(max_recv_token_num >= 0, + "MegaMoe requires max_recv_token_num >= 0."); + TORCH_CHECK(num_max_tokens_per_rank >= 0, + "MegaMoe requires num_max_tokens_per_rank >= 0."); + TORCH_CHECK(rank_num_per_server > 0, + "MegaMoe requires rank_num_per_server > 0."); + + auto output = at::empty_like(x); + auto expert_token_nums = + at::empty({local_expert_num}, x.options().dtype(at::kInt)); + + std::string comm_alg_copy = comm_alg; + char* comm_alg_ptr = comm_alg_copy.data(); + std::string activation_copy = activation; + char* activation_ptr = activation_copy.data(); + + // The generic bridge provides a testable record of the deployed custom ABI + // ordering while retaining EXEC_NPU_CMD's lvalue conversion contract. + auto context_arg = context; + auto x_arg = x; + auto topk_ids_arg = topk_ids; + auto topk_weights_arg = topk_weights; + auto weight1_arg = weight1; + auto weight2_arg = weight2; + auto weight_scales1_arg = weight_scales1; + auto weight_scales2_arg = weight_scales2; + auto bias1_arg = bias1; + auto bias2_arg = bias2; + auto x_active_mask_arg = x_active_mask; + auto execute = [&](auto&... arguments) { + EXEC_NPU_CMD(aclnnMegaMoe, arguments...); + }; + execute_mega_moe_acl_contract(execute, + context_arg, + x_arg, + topk_ids_arg, + topk_weights_arg, + weight1_arg, + weight2_arg, + weight_scales1_arg, + weight_scales2_arg, + bias1_arg, + bias2_arg, + x_active_mask_arg, + moe_expert_num, + ep_world_size, + ccl_buffer_size, + max_recv_token_num, + dispatch_quant_mode, + dispatch_quant_out_dtype, + combine_quant_mode, + comm_alg_ptr, + num_max_tokens_per_rank, + activation_ptr, + activation_clamp, + topo_type, + rank_num_per_server, + output, + expert_token_nums); + + return std::make_tuple(output, expert_token_nums); +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index dea67faf54..68ebd9c10c 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -265,6 +265,49 @@ std::tuple apply_npu_dispatch_ffn_combine( bool has_dispatch_ffn_combine(); +std::tuple apply_npu_mega_moe( + const torch::Tensor& context, + const torch::Tensor& x, + const torch::Tensor& topk_ids, + const torch::Tensor& topk_weights, + const torch::TensorList weight1, + const torch::TensorList weight2, + int64_t moe_expert_num, + int64_t ep_world_size, + int64_t ccl_buffer_size, + const std::optional& weight_scales1 = std::nullopt, + const std::optional& weight_scales2 = std::nullopt, + const std::optional& bias1 = std::nullopt, + const std::optional& bias2 = std::nullopt, + const std::optional& x_active_mask = std::nullopt, + int64_t max_recv_token_num = 0, + int64_t dispatch_quant_mode = 0, + int64_t combine_quant_mode = 0, + const std::string& comm_alg = "", + int64_t num_max_tokens_per_rank = 0, + const std::string& activation = "swiglu", + float activation_clamp = 3.402823466e+38F, + int64_t dispatch_quant_out_dtype = 28, + int64_t topo_type = 0, + int64_t rank_num_per_server = 2); + +struct MegaMoeOpApiProvenance { + bool compatible = false; + bool same_library = false; + std::string expected_library; + std::string workspace_library; + std::string execute_library; +}; + +MegaMoeOpApiProvenance validate_mega_moe_op_api_paths( + const std::string& expected_library, + const std::string& workspace_library, + const std::string& execute_library); + +MegaMoeOpApiProvenance inspect_mega_moe_op_api_provenance(); + +bool has_mega_moe(); + std::tuple apply_npu_dispatch_gmm_combine_decode( const torch::Tensor& x, const torch::Tensor& expert_ids, diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index b38addd7df..3c677b3586 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -1144,6 +1144,45 @@ bool has_dispatch_ffn_combine() { #endif } +std::tuple mega_moe(MegaMoeParams& params) { +#if defined(USE_NPU) + return npu::apply_npu_mega_moe(params.context, + params.x, + params.topk_ids, + params.topk_weights, + params.weight1, + params.weight2, + params.moe_expert_num, + params.ep_world_size, + params.ccl_buffer_size, + params.weight_scales1, + params.weight_scales2, + params.bias1, + params.bias2, + params.x_active_mask, + params.max_recv_token_num, + params.dispatch_quant_mode, + params.combine_quant_mode, + params.comm_alg, + params.num_max_tokens_per_rank, + params.activation, + params.activation_clamp, + params.dispatch_quant_out_dtype, + params.topo_type, + params.rank_num_per_server); +#else + NOT_IMPLEMENTED(); +#endif +} + +bool has_mega_moe() { +#if defined(USE_NPU) + return npu::has_mega_moe(); +#else + return false; +#endif +} + std::tuple dispatch_gmm_combine_decode( DispatchGmmCombineDecodeParams& params) { #if defined(USE_NPU) diff --git a/xllm/core/kernels/ops_api.h b/xllm/core/kernels/ops_api.h index 1ead17c396..1d2501fac0 100644 --- a/xllm/core/kernels/ops_api.h +++ b/xllm/core/kernels/ops_api.h @@ -141,6 +141,10 @@ std::tuple dispatch_ffn_combine( bool has_dispatch_ffn_combine(); +std::tuple mega_moe(MegaMoeParams& params); + +bool has_mega_moe(); + std::tuple dispatch_gmm_combine_decode( DispatchGmmCombineDecodeParams& params); diff --git a/xllm/core/kernels/param.h b/xllm/core/kernels/param.h index ec1b2ff182..d97fa43fd4 100644 --- a/xllm/core/kernels/param.h +++ b/xllm/core/kernels/param.h @@ -1411,6 +1411,43 @@ struct DispatchFFNCombineParams { std::optional expert_token_nums = std::nullopt; }; +// MegaMoe fuses EP dispatch, expert FFN, and EP combine. The policy layer +// admits only the verified A16W16 envelope; these defaults mirror the +// upstream torch extension and keep unsupported quantization paths disabled. +struct MegaMoeParams { + // Communication context from MegaMoeCommResource. Shape: [context_words], + // dtype int32, device NPU. + torch::Tensor context; + // Local input and global router results. + torch::Tensor x; + torch::Tensor topk_ids; + torch::Tensor topk_weights; + // Per-rank expert weights. W1[e]: [hidden, 2 * intermediate], + // W2[e]: [intermediate, hidden]. + torch::TensorList weight1; + torch::TensorList weight2; + int64_t moe_expert_num = 0; + int64_t ep_world_size = 0; + int64_t ccl_buffer_size = 0; + + std::optional weight_scales1 = std::nullopt; + std::optional weight_scales2 = std::nullopt; + std::optional bias1 = std::nullopt; + std::optional bias2 = std::nullopt; + std::optional x_active_mask = std::nullopt; + int64_t max_recv_token_num = 0; + int64_t dispatch_quant_mode = 0; + int64_t combine_quant_mode = 0; + std::string comm_alg; + int64_t num_max_tokens_per_rank = 0; + std::string activation = "swiglu"; + float activation_clamp = 3.402823466e+38F; + // 28 is the non-quantized sentinel used by the upstream torch bridge. + int64_t dispatch_quant_out_dtype = 28; + int64_t topo_type = 0; + int64_t rank_num_per_server = 2; +}; + struct DispatchGmmCombineDecodeParams { // Input hidden states. Shape: [num_tokens, hidden_size]. torch::Tensor x; diff --git a/xllm/core/layers/npu_torch/CMakeLists.txt b/xllm/core/layers/npu_torch/CMakeLists.txt index 0b09ff9751..5743f9548d 100755 --- a/xllm/core/layers/npu_torch/CMakeLists.txt +++ b/xllm/core/layers/npu_torch/CMakeLists.txt @@ -1,5 +1,26 @@ include(cc_library) +cc_library( + NAME + mega_moe_policy + HDRS + mega_moe_policy.h + SRCS + mega_moe_policy.cpp +) + +cc_library( + NAME + mega_moe_runtime + HDRS + mega_moe_runtime.h + SRCS + mega_moe_runtime.cpp + DEPS + torch +) + + cc_library( NAME npu_torch_layers @@ -39,6 +60,9 @@ cc_library( deepseek_v4_gate.cpp DEPS :common_layers + :mega_moe_policy + :mega_moe_comm_resource + :mega_moe_runtime :npu_layers :xllm_ops :triton_adapter diff --git a/xllm/core/layers/npu_torch/fused_moe.cpp b/xllm/core/layers/npu_torch/fused_moe.cpp index 4169124b0d..409e38a1b0 100644 --- a/xllm/core/layers/npu_torch/fused_moe.cpp +++ b/xllm/core/layers/npu_torch/fused_moe.cpp @@ -15,12 +15,16 @@ limitations under the License. #include "fused_moe.h" +#include #include +#include #include #include #include #include +#include +#include #include #include #include @@ -32,11 +36,17 @@ limitations under the License. #include #endif +#include "framework/config/execution_config.h" #include "framework/config/eplb_config.h" #include "framework/config/kernel_config.h" +#include "framework/config/scheduler_config.h" +#include "framework/parallel_state/mega_moe_comm_resource.h" #include "framework/parallel_state/parallel_state.h" +#include "kernels/npu/npu_ops_api.h" #include "kernels/ops_api.h" #include "layers/common/dp_utils.h" +#include "layers/npu_torch/mega_moe_policy.h" +#include "layers/npu_torch/mega_moe_runtime.h" #include "platform/device.h" #include "util/utils.h" @@ -63,6 +73,31 @@ torch::Tensor slice_expert_weights(const torch::Tensor& weight, .contiguous(); } +std::string read_mega_moe_vendor_version() { + const char* custom_opp_path = std::getenv("ASCEND_CUSTOM_OPP_PATH"); + if (custom_opp_path == nullptr || custom_opp_path[0] == '\0') { + return ""; + } + std::string root(custom_opp_path); + const size_t path_separator = root.find(':'); + if (path_separator != std::string::npos) { + root.resize(path_separator); + } + std::ifstream version_file(root + "/version.info"); + std::string line; + constexpr char kVersionPrefix[] = "custom_opp_compiler_version="; + while (std::getline(version_file, line)) { + if (line.rfind(kVersionPrefix, 0) == 0) { + return line.substr(sizeof(kVersionPrefix) - 1); + } + } + return ""; +} + +std::string mega_moe_rejection_text(MegaMoeRejectionReason reason) { + return std::to_string(static_cast(reason)); +} + std::optional resolve_moe_quant_method( const QuantArgs& quant_args, const StateDict& state_dict) { @@ -524,6 +559,103 @@ FusedMoEImpl::FusedMoEImpl(const ModelArgs& model_args, options_), false); } + initialize_mega_moe(); +} + +void FusedMoEImpl::initialize_mega_moe() { + const auto mode = parse_mega_moe_mode( + KernelConfig::get_instance().mega_moe_mode()); + CHECK(mode.has_value()) << "Invalid MegaMoe mode."; + if (mode.value() == MegaMoeMode::OFF) { + return; + } + + const int64_t weight_budget_limit = + KernelConfig::get_instance().mega_moe_weight_cache_budget_bytes(); + const int64_t estimated_weight_bytes = + estimate_mega_moe_weight_cache_bytes(w13_, w2_); + const auto weight_budget = inspect_mega_moe_weight_budget( + estimated_weight_bytes, weight_budget_limit); + if (!weight_budget.allowed) { + if (mode.value() == MegaMoeMode::AUTO) { + LOG(WARNING) << "MegaMoe auto fallback before HCCL: weight cache " + << "budget exceeded, estimated=" + << weight_budget.estimated_bytes + << ", current=" << weight_budget.current_bytes + << ", limit=" << weight_budget.limit_bytes; + return; + } + CHECK(weight_budget.allowed) + << "MegaMoe forced-on weight cache budget exceeded: estimated=" + << weight_budget.estimated_bytes + << ", current=" << weight_budget.current_bytes + << ", limit=" << weight_budget.limit_bytes; + } + if (mode.value() == MegaMoeMode::AUTO) { + LOG(WARNING) << "MegaMoe auto mode is not enabled by this validation " + "gate; falling back to legacy before HCCL access."; + return; + } + + const char* soc_name = aclrtGetSocName(); + const auto comm_symbols = probe_mega_moe_comm_symbols(); + const auto op_provenance = + xllm::kernel::npu::inspect_mega_moe_op_api_provenance(); + const bool op_symbols = op_provenance.compatible; + const bool unquantized = + quant_args_.quant_method().empty() && quant_args_.quant_descs().empty(); + const bool swiglu = + is_gated_ && (hidden_act_ == "silu" || hidden_act_ == "swiglu"); + MegaMoeCapability capability; + capability.soc_name = soc_name == nullptr ? "" : soc_name; + capability.cann_version = CANN_VERSION_STR; + capability.vendor_version = read_mega_moe_vendor_version(); + capability.dtype = options_.dtype() == torch::kBFloat16 + ? MegaMoeDType::BFLOAT16 + : MegaMoeDType::OTHER; + capability.activation = swiglu ? "swiglu" : hidden_act_; + capability.hidden_size = hidden_size_; + capability.intermediate_size = local_intermediate_size_; + capability.num_experts = num_total_experts_; + capability.topk = topk_; + capability.ep_size = parallel_args_.ep_size(); + capability.tp_size = tp_pg_->world_size(); + capability.dp_size = parallel_args_.dp_size(); + capability.graph_enabled = ExecutionConfig::get_instance().enable_graph(); + capability.quantization_enabled = !unquantized; + capability.workspace_symbol_ready = op_symbols; + capability.execute_symbol_ready = op_symbols; + capability.comm_context_ready = comm_symbols.available && + parallel_args_.moe_ep_group_ != nullptr; + capability.fused_mc2_enabled = fused_mc2_mode() != 0; + capability.collectives_started = false; + capability.weight_budget_ready = weight_budget.allowed; + + const MegaMoeDecision decision = decide_mega_moe(mode.value(), capability); + CHECK(!decision.fatal) + << "MegaMoe forced-on preflight failed, reason=" + << mega_moe_rejection_text(decision.reason) + << ", expected_op_library=" << op_provenance.expected_library + << ", workspace_library=" << op_provenance.workspace_library + << ", execute_library=" << op_provenance.execute_library + << ". Source CANN 9.1 and exactly one validated custom vendor; set " + "ASCEND_CUSTOM_OPP_PATH to that vendor root."; + CHECK(decision.use_mega_moe); + + ProcessGroup* ep_group = parallel_args_.moe_ep_group_; + MegaMoeCommSpec comm_spec; + comm_spec.group_name = ep_group->hccl_comm_name(/*init_comm=*/true); + comm_spec.hccl_comm = ep_group->hccl_comm(); + comm_spec.ep_world_size = ep_group->world_size(); + comm_spec.device_index = options_.device().index(); + comm_spec.max_num_tokens_per_rank = + SchedulerConfig::get_instance().max_tokens_per_batch(); + mega_moe_comm_resource_ = ep_group->acquire_mega_moe_comm_resource(comm_spec); + mega_moe_enabled_ = true; + LOG(INFO) << "MegaMoe forced-on path enabled for EP=" + << comm_spec.ep_world_size + << ", max_tokens_per_rank=" + << comm_spec.max_num_tokens_per_rank; } void FusedMoEImpl::validate_resolved_quant_method() const { @@ -742,10 +874,10 @@ void FusedMoEImpl::ensure_group_gemm_weight_layout(torch::Tensor& weight, << weight.sizes() << ", expected dim2=" << output_dim; } -torch::Tensor FusedMoEImpl::select_experts( +std::pair +FusedMoEImpl::select_global_experts( const torch::Tensor& hidden_states_2d, - const torch::Tensor& router_logits_2d, - SelectedExpertInfo& selected_expert_info) { + const torch::Tensor& router_logits_2d) { torch::Tensor topk_weights; torch::Tensor topk_ids; std::optional e_score_correction_bias = std::nullopt; @@ -805,6 +937,16 @@ torch::Tensor FusedMoEImpl::select_experts( } } + return std::make_pair(topk_weights.contiguous(), topk_ids.contiguous()); +} + +torch::Tensor FusedMoEImpl::select_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d, + SelectedExpertInfo& selected_expert_info) { + auto [topk_weights, topk_ids] = + select_global_experts(hidden_states_2d, router_logits_2d); + const int64_t local_expert_start = start_expert_id_; const int64_t local_expert_end = start_expert_id_ + num_experts_per_rank_; if (parallel_args_.ep_size() > 1) { @@ -849,10 +991,93 @@ torch::Tensor FusedMoEImpl::select_experts( return expand_hidden_states; } +void FusedMoEImpl::ensure_mega_moe_weight_cache( + bool require_complete_weights) { + const auto action = plan_mega_moe_weight_cache( + mega_moe_enabled_, + mega_moe_weight_cache_.ready(), + w13_is_loaded_, + w2_is_loaded_, + require_complete_weights); + if (action == MegaMoeWeightCacheAction::SKIP || + action == MegaMoeWeightCacheAction::WAIT_FOR_WEIGHTS) { + return; + } + CHECK(action != MegaMoeWeightCacheAction::FAIL_MISSING_WEIGHTS) + << "MegaMoe weights must be loaded before first execution."; + CHECK(action == MegaMoeWeightCacheAction::BUILD); + MegaMoeWeightBudgetStatus budget_status; + auto weight_cache = try_build_mega_moe_weight_cache( + w13_, + w2_, + hidden_size_, + local_intermediate_size_, + KernelConfig::get_instance().mega_moe_weight_cache_budget_bytes(), + &budget_status); + CHECK(weight_cache.has_value()) + << "MegaMoe forced-on weight cache reservation failed before " + "contiguous allocation: estimated=" + << budget_status.estimated_bytes + << ", current=" << budget_status.current_bytes + << ", limit=" << budget_status.limit_bytes; + mega_moe_weight_cache_ = std::move(weight_cache.value()); + LOG(INFO) << "MegaMoe materialized independent contiguous weight caches: " + << mega_moe_weight_cache_.memory_bytes + << " bytes; process_total=" + << current_mega_moe_weight_cache_bytes() << "."; +} + +torch::Tensor FusedMoEImpl::forward_mega_moe( + const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + const std::optional& shared_output) { + CHECK(mega_moe_enabled_); + auto comm_resource = mega_moe_comm_resource_.lock(); + CHECK(comm_resource != nullptr) + << "MegaMoe communication resource expired before collective launch."; + + ensure_mega_moe_weight_cache(/*require_complete_weights=*/true); + const auto hidden_states_shape = hidden_states.sizes(); + auto input_2d = + hidden_states.reshape({-1, hidden_states.size(-1)}).contiguous(); + auto router_logits_2d = + router_logits.reshape({-1, router_logits.size(-1)}); + auto [topk_weights, topk_ids] = + select_global_experts(input_2d, router_logits_2d); + topk_weights = prepare_mega_moe_topk_weights(topk_weights); + + xllm::kernel::MegaMoeParams params; + params.context = comm_resource->context_tensor(); + params.x = input_2d; + params.topk_ids = topk_ids; + params.topk_weights = topk_weights; + params.weight1 = torch::TensorList(mega_moe_weight_cache_.weight1); + params.weight2 = torch::TensorList(mega_moe_weight_cache_.weight2); + params.moe_expert_num = num_total_experts_; + params.ep_world_size = parallel_args_.ep_size(); + params.ccl_buffer_size = comm_resource->ccl_buffer_size(); + params.num_max_tokens_per_rank = + comm_resource->max_num_tokens_per_rank(); + params.activation = "swiglu"; + + torch::Tensor routed_output; + torch::Tensor expert_token_nums; + std::tie(routed_output, expert_token_nums) = xllm::kernel::mega_moe(params); + (void)expert_token_nums; + auto output = routed_output.reshape(hidden_states_shape); + if (shared_output.has_value()) { + output = output + shared_output.value(); + } + return output; +} + torch::Tensor FusedMoEImpl::forward_expert( const torch::Tensor& hidden_states, const torch::Tensor& router_logits, const std::optional& shared_output) { + if (mega_moe_enabled_) { + return forward_mega_moe(hidden_states, router_logits, shared_output); + } // prepare the parameters for MoE computation torch::IntArrayRef hidden_states_shape = hidden_states.sizes(); torch::ScalarType hidden_states_dtype = hidden_states.dtype().toScalarType(); @@ -1962,6 +2187,7 @@ void FusedMoEImpl::load_state_dict(const StateDict& state_dict) { } load_experts(state_dict.get_dict_with_prefix("experts.")); preprocess_w4a8_dynamic_weights(); + ensure_mega_moe_weight_cache(/*require_complete_weights=*/false); } } // namespace layer diff --git a/xllm/core/layers/npu_torch/fused_moe.h b/xllm/core/layers/npu_torch/fused_moe.h index e21b9bf50d..77ee0179b0 100644 --- a/xllm/core/layers/npu_torch/fused_moe.h +++ b/xllm/core/layers/npu_torch/fused_moe.h @@ -17,6 +17,7 @@ limitations under the License. #include +#include #include #include #include @@ -31,8 +32,13 @@ limitations under the License. #include "layers/common/dense_mlp.h" #include "layers/common/fused_moe_base.h" #include "layers/common/linear.h" +#include "layers/npu_torch/mega_moe_policy.h" +#include "layers/npu_torch/mega_moe_runtime.h" namespace xllm { + +class MegaMoeCommResource; + namespace layer { class FusedMoEImpl : public torch::nn::Module { @@ -67,12 +73,27 @@ class FusedMoEImpl : public torch::nn::Module { std::optional input_scale; }; + // Computes router choices in the global expert-id space. MegaMoe consumes + // these values before the legacy per-rank expert mask is applied. + std::pair select_global_experts( + const torch::Tensor& hidden_states_2d, + const torch::Tensor& router_logits_2d); // initial steps for MoE computation, select the experts for each token torch::Tensor select_experts(const torch::Tensor& hidden_states_2d, const torch::Tensor& router_logits_2d, SelectedExpertInfo& selected_expert_info); - private: + void initialize_mega_moe(); + void ensure_mega_moe_weight_cache(bool require_complete_weights); + torch::Tensor forward_mega_moe( + const torch::Tensor& hidden_states, + const torch::Tensor& router_logits, + const std::optional& shared_output); + + bool mega_moe_enabled_ = false; + std::weak_ptr mega_moe_comm_resource_; + MegaMoeWeightCache mega_moe_weight_cache_; + int64_t num_total_experts_; int64_t topk_; int64_t num_expert_group_; diff --git a/xllm/core/layers/npu_torch/mega_moe_policy.cpp b/xllm/core/layers/npu_torch/mega_moe_policy.cpp new file mode 100644 index 0000000000..7a2f9e1641 --- /dev/null +++ b/xllm/core/layers/npu_torch/mega_moe_policy.cpp @@ -0,0 +1,155 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "layers/npu_torch/mega_moe_policy.h" + +namespace xllm { +namespace layer { +namespace { + +constexpr char kSupportedSocPrefix[] = "Ascend910_93"; +constexpr char kSupportedVersion[] = "9.1.0-beta.3"; + +MegaMoeRejectionReason find_rejection_reason( + const MegaMoeCapability& capability) { + if (capability.fused_mc2_enabled) { + return MegaMoeRejectionReason::FUSED_MC2_CONFLICT; + } + if (capability.soc_name.rfind(kSupportedSocPrefix, 0) != 0) { + return MegaMoeRejectionReason::UNSUPPORTED_SOC; + } + if (capability.cann_version != kSupportedVersion) { + return MegaMoeRejectionReason::UNSUPPORTED_CANN_VERSION; + } + if (capability.vendor_version != kSupportedVersion) { + return MegaMoeRejectionReason::UNSUPPORTED_VENDOR_VERSION; + } + if (!capability.workspace_symbol_ready || + !capability.execute_symbol_ready) { + return MegaMoeRejectionReason::OP_SYMBOLS_UNAVAILABLE; + } + if (!capability.comm_context_ready) { + return MegaMoeRejectionReason::COMM_CONTEXT_UNAVAILABLE; + } + if (capability.dtype != MegaMoeDType::BFLOAT16) { + return MegaMoeRejectionReason::UNSUPPORTED_DTYPE; + } + if (capability.activation != "swiglu") { + return MegaMoeRejectionReason::UNSUPPORTED_ACTIVATION; + } + if (capability.hidden_size != 2048 || + capability.intermediate_size != 512 || + capability.num_experts != 256 || capability.topk != 8) { + return MegaMoeRejectionReason::UNSUPPORTED_SHAPE; + } + if (capability.ep_size <= 0 || + capability.num_experts % capability.ep_size != 0) { + return MegaMoeRejectionReason::UNSUPPORTED_EP; + } + if (capability.tp_size != 1) { + return MegaMoeRejectionReason::UNSUPPORTED_TP; + } + if (capability.graph_enabled) { + return MegaMoeRejectionReason::GRAPH_ENABLED; + } + if (capability.quantization_enabled) { + return MegaMoeRejectionReason::QUANTIZATION_ENABLED; + } + if (!capability.weight_budget_ready) { + return MegaMoeRejectionReason::WEIGHT_BUDGET_EXCEEDED; + } + return MegaMoeRejectionReason::NONE; +} + +} // namespace + +std::optional parse_mega_moe_mode(const std::string& value) { + if (value == "off") { + return MegaMoeMode::OFF; + } + if (value == "auto") { + return MegaMoeMode::AUTO; + } + if (value == "on") { + return MegaMoeMode::ON; + } + return std::nullopt; +} + +MegaMoeDecision decide_mega_moe( + MegaMoeMode mode, + const MegaMoeCapability& capability) { + if (mode == MegaMoeMode::OFF) { + return MegaMoeDecision{ + .use_mega_moe = false, + .fatal = false, + .reason = MegaMoeRejectionReason::DISABLED}; + } + + const MegaMoeRejectionReason rejection_reason = + find_rejection_reason(capability); + if (rejection_reason == MegaMoeRejectionReason::NONE) { + return MegaMoeDecision{ + .use_mega_moe = true, + .fatal = false, + .reason = MegaMoeRejectionReason::NONE}; + } + + // MegaMoe and fused_mc2 both own the routed-expert communication sequence. + // Falling back silently would make precedence configuration-dependent. + if (rejection_reason == MegaMoeRejectionReason::FUSED_MC2_CONFLICT) { + return MegaMoeDecision{.use_mega_moe = false, + .fatal = true, + .reason = rejection_reason}; + } + + if (mode == MegaMoeMode::AUTO && !capability.collectives_started) { + return MegaMoeDecision{ + .use_mega_moe = false, + .fatal = false, + .reason = rejection_reason}; + } + + const MegaMoeRejectionReason fatal_reason = + mode == MegaMoeMode::AUTO + ? MegaMoeRejectionReason::COLLECTIVES_ALREADY_STARTED + : rejection_reason; + return MegaMoeDecision{ + .use_mega_moe = false, + .fatal = true, + .reason = fatal_reason}; +} + +MegaMoeExecutionContract make_mega_moe_execution_contract( + const MegaMoeDecision& decision, + int64_t ep_size, + bool has_shared_expert) { + const bool distributed_ep = ep_size > 1; + if (decision.use_mega_moe) { + return MegaMoeExecutionContract{ + .preserve_global_topk = true, + .apply_local_expert_mask = false, + .reduce_routed_output_across_ep = false, + .shared_output_additions = has_shared_expert ? 1 : 0}; + } + return MegaMoeExecutionContract{ + .preserve_global_topk = false, + .apply_local_expert_mask = distributed_ep, + .reduce_routed_output_across_ep = distributed_ep, + .shared_output_additions = has_shared_expert ? 1 : 0}; +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/npu_torch/mega_moe_policy.h b/xllm/core/layers/npu_torch/mega_moe_policy.h new file mode 100644 index 0000000000..0da49d875f --- /dev/null +++ b/xllm/core/layers/npu_torch/mega_moe_policy.h @@ -0,0 +1,104 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include +#include +#include + +namespace xllm { +namespace layer { + +enum class MegaMoeMode : int8_t { + OFF = 0, + AUTO = 1, + ON = 2, +}; + +enum class MegaMoeDType : int8_t { + BFLOAT16 = 0, + OTHER = 1, +}; + +enum class MegaMoeRejectionReason : int8_t { + NONE = 0, + DISABLED = 1, + FUSED_MC2_CONFLICT = 2, + COLLECTIVES_ALREADY_STARTED = 3, + UNSUPPORTED_SOC = 4, + UNSUPPORTED_CANN_VERSION = 5, + UNSUPPORTED_VENDOR_VERSION = 6, + OP_SYMBOLS_UNAVAILABLE = 7, + COMM_CONTEXT_UNAVAILABLE = 8, + UNSUPPORTED_DTYPE = 9, + UNSUPPORTED_ACTIVATION = 10, + UNSUPPORTED_SHAPE = 11, + UNSUPPORTED_EP = 12, + UNSUPPORTED_TP = 13, + UNSUPPORTED_DP = 14, + GRAPH_ENABLED = 15, + QUANTIZATION_ENABLED = 16, + WEIGHT_BUDGET_EXCEEDED = 17, +}; + +struct MegaMoeCapability { + std::string soc_name; + std::string cann_version; + std::string vendor_version; + MegaMoeDType dtype = MegaMoeDType::OTHER; + std::string activation; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_experts = 0; + int64_t topk = 0; + int64_t ep_size = 0; + int64_t tp_size = 0; + int64_t dp_size = 0; + bool graph_enabled = false; + bool quantization_enabled = false; + bool workspace_symbol_ready = false; + bool execute_symbol_ready = false; + bool comm_context_ready = false; + bool fused_mc2_enabled = false; + bool collectives_started = false; + bool weight_budget_ready = true; +}; + +struct MegaMoeDecision { + bool use_mega_moe = false; + bool fatal = false; + MegaMoeRejectionReason reason = MegaMoeRejectionReason::NONE; +}; + +struct MegaMoeExecutionContract { + bool preserve_global_topk = false; + bool apply_local_expert_mask = true; + bool reduce_routed_output_across_ep = true; + int32_t shared_output_additions = 0; +}; + +std::optional parse_mega_moe_mode(const std::string& value); + +MegaMoeDecision decide_mega_moe(MegaMoeMode mode, + const MegaMoeCapability& capability); + +MegaMoeExecutionContract make_mega_moe_execution_contract( + const MegaMoeDecision& decision, + int64_t ep_size, + bool has_shared_expert); + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/layers/npu_torch/mega_moe_runtime.cpp b/xllm/core/layers/npu_torch/mega_moe_runtime.cpp new file mode 100644 index 0000000000..14a5dfcbeb --- /dev/null +++ b/xllm/core/layers/npu_torch/mega_moe_runtime.cpp @@ -0,0 +1,245 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include "layers/npu_torch/mega_moe_runtime.h" + +#include +#include +#include +#include + +namespace xllm::layer { +namespace { + +std::atomic g_weight_cache_bytes{0}; + +class BudgetReservation final { + public: + BudgetReservation(int64_t bytes, std::function on_release) + : bytes_(bytes), on_release_(std::move(on_release)) {} + ~BudgetReservation() { + g_weight_cache_bytes.fetch_sub(bytes_, std::memory_order_acq_rel); + if (on_release_) { + on_release_(); + } + } + + private: + int64_t bytes_; + std::function on_release_; +}; + +class TensorFieldsReleaseProbe final { + public: + explicit TensorFieldsReleaseProbe(std::function on_release) + : on_release_(std::move(on_release)) {} + ~TensorFieldsReleaseProbe() { + if (on_release_) { + on_release_(); + } + } + + private: + std::function on_release_; +}; + +void validate_canonical_weights( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + int64_t hidden_size, + int64_t intermediate_size) { + TORCH_CHECK(canonical_w13.defined() && canonical_w2.defined(), + "MegaMoe canonical weights must be defined."); + TORCH_CHECK(canonical_w13.dim() == 3 && canonical_w2.dim() == 3, + "MegaMoe canonical weights must be 3D."); + TORCH_CHECK(canonical_w13.scalar_type() == at::kBFloat16 && + canonical_w2.scalar_type() == at::kBFloat16, + "MegaMoe weight cache requires bf16 canonical weights."); + TORCH_CHECK(canonical_w13.size(0) == canonical_w2.size(0), + "MegaMoe canonical expert count mismatch: ", + canonical_w13.size(0), + " vs ", + canonical_w2.size(0)); + TORCH_CHECK(canonical_w13.size(1) == 2 * intermediate_size && + canonical_w13.size(2) == hidden_size, + "MegaMoe canonical W13 must be [expert, 2I, H], got ", + canonical_w13.sizes()); + TORCH_CHECK(canonical_w2.size(1) == hidden_size && + canonical_w2.size(2) == intermediate_size, + "MegaMoe canonical W2 must be [expert, H, I], got ", + canonical_w2.sizes()); +} + +std::shared_ptr try_reserve_budget( + int64_t estimated_bytes, + int64_t limit_bytes, + MegaMoeWeightBudgetStatus& status, + std::function on_release) { + int64_t current = g_weight_cache_bytes.load(std::memory_order_acquire); + while (true) { + status = inspect_mega_moe_weight_budget(estimated_bytes, limit_bytes); + status.current_bytes = current; + status.allowed = estimated_bytes > 0 && limit_bytes >= estimated_bytes && + current <= limit_bytes - estimated_bytes; + if (!status.allowed) { + return nullptr; + } + if (g_weight_cache_bytes.compare_exchange_weak( + current, + current + estimated_bytes, + std::memory_order_acq_rel, + std::memory_order_acquire)) { + break; + } + } + + try { + return std::static_pointer_cast( + std::make_shared(estimated_bytes, + std::move(on_release))); + } catch (...) { + g_weight_cache_bytes.fetch_sub(estimated_bytes, + std::memory_order_acq_rel); + throw; + } +} + +void materialize_weight_cache(const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + MegaMoeWeightCache& cache) { + // contiguous() materializes independent storage; canonical registered + // parameters remain in the legacy loader layout and stay readable. + cache.weight1_storage = canonical_w13.transpose(1, 2).contiguous(); + cache.weight2_storage = canonical_w2.transpose(1, 2).contiguous(); + const int64_t expert_count = canonical_w13.size(0); + cache.weight1.reserve(static_cast(expert_count)); + cache.weight2.reserve(static_cast(expert_count)); + for (int64_t expert = 0; expert < expert_count; ++expert) { + cache.weight1.push_back(cache.weight1_storage.select(0, expert)); + cache.weight2.push_back(cache.weight2_storage.select(0, expert)); + } +} + +} // namespace + +MegaMoeWeightCacheAction plan_mega_moe_weight_cache( + bool mega_moe_enabled, + bool cache_ready, + bool w13_loaded, + bool w2_loaded, + bool require_complete_weights) { + if (!mega_moe_enabled || cache_ready) { + return MegaMoeWeightCacheAction::SKIP; + } + if (!w13_loaded || !w2_loaded) { + return require_complete_weights + ? MegaMoeWeightCacheAction::FAIL_MISSING_WEIGHTS + : MegaMoeWeightCacheAction::WAIT_FOR_WEIGHTS; + } + return MegaMoeWeightCacheAction::BUILD; +} + +torch::Tensor prepare_mega_moe_topk_weights( + const torch::Tensor& topk_weights) { + TORCH_CHECK(topk_weights.defined(), + "MegaMoe top-k weights must be defined."); + TORCH_CHECK(topk_weights.dim() == 2, + "MegaMoe top-k weights must be 2D."); + return topk_weights.to(torch::kFloat32).contiguous(); +} + +int64_t estimate_mega_moe_weight_cache_bytes( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2) { + TORCH_CHECK(canonical_w13.defined() && canonical_w2.defined(), + "MegaMoe canonical weights must be defined for estimation."); + return canonical_w13.numel() * + static_cast(canonical_w13.element_size()) + + canonical_w2.numel() * + static_cast(canonical_w2.element_size()); +} + +int64_t current_mega_moe_weight_cache_bytes() { + return g_weight_cache_bytes.load(std::memory_order_acquire); +} + +MegaMoeWeightBudgetStatus inspect_mega_moe_weight_budget( + int64_t estimated_bytes, + int64_t limit_bytes) { + const int64_t current = current_mega_moe_weight_cache_bytes(); + const bool allowed = + estimated_bytes > 0 && limit_bytes >= estimated_bytes && + current <= limit_bytes - estimated_bytes; + return {allowed, estimated_bytes, current, limit_bytes}; +} + +std::optional try_build_mega_moe_weight_cache( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + int64_t hidden_size, + int64_t intermediate_size, + int64_t budget_limit_bytes, + MegaMoeWeightBudgetStatus* status, + MegaMoeWeightCacheDestructionObserver* observer) { + validate_canonical_weights( + canonical_w13, canonical_w2, hidden_size, intermediate_size); + const int64_t estimated_bytes = + estimate_mega_moe_weight_cache_bytes(canonical_w13, canonical_w2); + MegaMoeWeightBudgetStatus reservation_status; + std::function on_reservation_released; + if (observer != nullptr) { + on_reservation_released = observer->on_reservation_released; + } + std::shared_ptr reservation = try_reserve_budget( + estimated_bytes, + budget_limit_bytes, + reservation_status, + std::move(on_reservation_released)); + if (status != nullptr) { + *status = reservation_status; + } + if (reservation == nullptr) { + return std::nullopt; + } + + MegaMoeWeightCache cache; + cache.memory_bytes = estimated_bytes; + cache.budget_reservation_ = std::move(reservation); + if (observer != nullptr && observer->on_tensor_fields_released) { + cache.tensor_fields_release_probe_ = std::static_pointer_cast( + std::make_shared( + observer->on_tensor_fields_released)); + } + materialize_weight_cache(canonical_w13, canonical_w2, cache); + return cache; +} + +MegaMoeWeightCache build_mega_moe_weight_cache( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + int64_t hidden_size, + int64_t intermediate_size) { + auto cache = try_build_mega_moe_weight_cache( + canonical_w13, + canonical_w2, + hidden_size, + intermediate_size, + std::numeric_limits::max()); + TORCH_CHECK(cache.has_value(), + "MegaMoe weight cache accounting overflow."); + return std::move(cache.value()); +} + +} // namespace xllm::layer diff --git a/xllm/core/layers/npu_torch/mega_moe_runtime.h b/xllm/core/layers/npu_torch/mega_moe_runtime.h new file mode 100644 index 0000000000..cd1bb9a9ae --- /dev/null +++ b/xllm/core/layers/npu_torch/mega_moe_runtime.h @@ -0,0 +1,121 @@ +/* Copyright 2026 The xLLM Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include +#include + +namespace xllm::layer { + +struct MegaMoeWeightBudgetStatus { + bool allowed = false; + int64_t estimated_bytes = 0; + int64_t current_bytes = 0; + int64_t limit_bytes = 0; +}; + +enum class MegaMoeWeightCacheAction : int32_t { + SKIP = 0, + WAIT_FOR_WEIGHTS, + BUILD, + FAIL_MISSING_WEIGHTS, +}; + +MegaMoeWeightCacheAction plan_mega_moe_weight_cache( + bool mega_moe_enabled, + bool cache_ready, + bool w13_loaded, + bool w2_loaded, + bool require_complete_weights); + +torch::Tensor prepare_mega_moe_topk_weights( + const torch::Tensor& topk_weights); + +struct MegaMoeWeightCacheDestructionObserver { + std::function on_tensor_fields_released; + std::function on_reservation_released; +}; + +struct MegaMoeWeightCache { + MegaMoeWeightCache() = default; + MegaMoeWeightCache(const MegaMoeWeightCache&) = delete; + MegaMoeWeightCache& operator=(const MegaMoeWeightCache&) = delete; + MegaMoeWeightCache(MegaMoeWeightCache&&) noexcept = default; + MegaMoeWeightCache& operator=(MegaMoeWeightCache&&) noexcept = default; + + private: + friend std::optional + try_build_mega_moe_weight_cache( + const torch::Tensor&, + const torch::Tensor&, + int64_t, + int64_t, + int64_t, + MegaMoeWeightBudgetStatus*, + MegaMoeWeightCacheDestructionObserver*); + // Members are destroyed in reverse declaration order. Keep the accounting + // reservation before the release probe and tensor fields so device storage + // is released before the process-wide budget is returned. + std::shared_ptr budget_reservation_; + std::shared_ptr tensor_fields_release_probe_; + + public: + torch::Tensor weight1_storage; + torch::Tensor weight2_storage; + std::vector weight1; + std::vector weight2; + int64_t memory_bytes = 0; + + bool ready() const { + return weight1_storage.defined() && weight2_storage.defined() && + !weight1.empty() && weight1.size() == weight2.size(); + } + +}; + +MegaMoeWeightCache build_mega_moe_weight_cache( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + int64_t hidden_size, + int64_t intermediate_size); + +int64_t estimate_mega_moe_weight_cache_bytes( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2); + +int64_t current_mega_moe_weight_cache_bytes(); + +MegaMoeWeightBudgetStatus inspect_mega_moe_weight_budget( + int64_t estimated_bytes, + int64_t limit_bytes); + +// Reserves the process-wide budget atomically before contiguous allocation. +// A denied reservation returns nullopt without allocating either weight copy. +std::optional try_build_mega_moe_weight_cache( + const torch::Tensor& canonical_w13, + const torch::Tensor& canonical_w2, + int64_t hidden_size, + int64_t intermediate_size, + int64_t budget_limit_bytes, + MegaMoeWeightBudgetStatus* status = nullptr, + MegaMoeWeightCacheDestructionObserver* observer = nullptr); + +} // namespace xllm::layer diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 6390495833..51d6f86438 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -203,9 +203,22 @@ void prepare_input_params_for_linear_attention(ModelInputParams& input_params) { input_params.parallel.query_start_loc[i + 1] = input_params.parallel.query_start_loc[i] + seq_len; } + const auto& kv_cache_tokens_nums = + input_params.attention.device.kv_cache_tokens_nums; + if (!kv_cache_tokens_nums.defined() || + kv_cache_tokens_nums.numel() == 0) { + CHECK_EQ(input_params.meta.num_sequences, 0) + << "kv_cache_tokens_nums may only be absent on an empty DP/EP shard"; + // Empty DP/EP shards still execute one synthetic token so every rank + // participates in collectives. Keep the linear-attention metadata + // coherent with that token without turning it into a scheduled sequence. + input_params.parallel.query_start_loc = {0, 1}; + input_params.parallel.has_initial_state = {0}; + return; + } torch::Tensor has_initial_state_tensor = - input_params.attention.device.kv_cache_tokens_nums > 0; + kv_cache_tokens_nums > 0; torch::Tensor has_initial_state_int64 = has_initial_state_tensor.contiguous() .view({-1}) .to(torch::kCPU) @@ -942,13 +955,26 @@ void WorkerImpl::prepare_work_before_execute_on_stream( bool empty_shard = input_params.meta.num_sequences == 0 && (!processed_input.token_ids.defined() || processed_input.token_ids.numel() == 0); - const bool need_fake_input_for_empty_shard = - empty_shard && !input_params.meta.batch_forward_type.is_empty() && + const bool uses_fake_inputs_for_empty_shards = + !input_params.meta.batch_forward_type.is_empty() && ((context_.get_parallel_args().cp_size() > 1 && !Platform::uses_model_cp_partition()) || (context_.get_parallel_args().dp_size() > 1 || context_.get_parallel_args().ep_size() > 1 || !context_.get_parallel_args().mapping_data().empty())); + const bool need_fake_input_for_empty_shard = + empty_shard && uses_fake_inputs_for_empty_shards; + if (uses_fake_inputs_for_empty_shards) { + // The runtime materializes one synthetic token for every empty shard. + // All ranks must therefore use the same collective tensor lengths, + // including active ranks that do not enter the fake-input branch below. + for (int32_t& token_num : + input_params.parallel.dp_global_token_nums) { + if (token_num == 0) { + token_num = 1; + } + } + } if (need_fake_input_for_empty_shard) { auto token_options = processed_input.token_ids.defined() ? processed_input.token_ids.options() @@ -960,6 +986,10 @@ void WorkerImpl::prepare_work_before_execute_on_stream( torch::ones({1}, token_options.device(device_)); processed_input.positions = torch::zeros({1}, position_options.device(device_)); + input_params.embedding.linear_state_ids = {0}; + input_params.attention.host.q_seq_lens = {1}; + input_params.attention.host.q_cu_seq_lens = {0, 1}; + input_params.attention.host.kv_seq_lens = {1}; empty_shard = false; } if (empty_shard) { @@ -1008,7 +1038,8 @@ void WorkerImpl::prepare_work_before_execute_on_stream( // Defer the slot-restore copy to step_for_schedule_overlap (worker // thread, on compute_stream_) so stream ordering between chunk N-1 // writes and chunk N restore is automatic. - if (!enable_schedule_overlap()) { + if (!enable_schedule_overlap() && + !need_fake_input_for_empty_shard) { restore_linear_state_slots( kv_caches_, processed_input.input_params.linear_state_cache_ops, diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index be8b0a5e9a..311000f9a3 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -96,7 +96,8 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { layer::AttentionMetadataBuilder::build( input_params, model_args_.enable_mla(), - build_attention_mask(input_params)); + build_attention_mask(input_params), + device_); torch::Tensor h; if (input_params.embedding.input_embedding.defined()) { h = input_params.embedding.input_embedding;