From 2c1850eb9f574eee804447e19920641be9235a14 Mon Sep 17 00:00:00 2001 From: Cooofish <2716856597@qq.com> Date: Tue, 7 Jul 2026 09:58:07 +0800 Subject: [PATCH 1/6] feat: add FlashComm1 and MMRS fusion support --- xllm/core/common/CMakeLists.txt | 3 + xllm/core/common/flash_comm1_context.cpp | 356 ++++++++++++++++++ xllm/core/common/flash_comm1_context.h | 103 +++++ xllm/core/common/global_flags.h | 11 + xllm/core/common/options.cpp | 6 + xllm/core/common/options.h | 11 + xllm/core/distributed_runtime/master.cpp | 18 + xllm/core/framework/model_context.cpp | 1 + xllm/core/framework/model_context.h | 10 + xllm/core/kernels/npu/CMakeLists.txt | 1 + xllm/core/kernels/npu/matmul.cpp | 1 - .../kernels/npu/matmul_reduce_scatter.cpp | 190 ++++++++++ xllm/core/kernels/npu/npu_ops_api.h | 15 + xllm/core/kernels/ops_api.cpp | 20 + xllm/core/kernels/ops_api.h | 2 + xllm/core/kernels/param.h | 24 ++ xllm/core/layers/common/dense_mlp.cpp | 49 ++- xllm/core/layers/common/dense_mlp.h | 4 +- xllm/core/layers/common/linear.cpp | 264 +++++++++++++ xllm/core/layers/common/linear.h | 8 + .../npu_torch/qwen3_gated_delta_net_base.cpp | 26 +- .../npu_torch/qwen3_gated_delta_net_base.h | 7 + .../layers/npu_torch/qwen3_next_attention.cpp | 37 +- .../layers/npu_torch/qwen3_next_attention.h | 16 +- .../qwen3_next_hybrid_decoder_layer_base.cpp | 50 ++- .../qwen3_next_hybrid_decoder_layer_base.h | 42 ++- xllm/core/runtime/llm_worker_impl.cpp | 10 +- xllm/core/runtime/options.h | 11 + xllm/core/runtime/worker_impl.cpp | 11 + xllm/models/llm/qwen3_next_hybrid_base.h | 25 +- xllm/pybind/bind.cpp | 7 + xllm/xllm.cpp | 6 + 32 files changed, 1296 insertions(+), 49 deletions(-) create mode 100644 xllm/core/common/flash_comm1_context.cpp create mode 100644 xllm/core/common/flash_comm1_context.h create mode 100644 xllm/core/kernels/npu/matmul_reduce_scatter.cpp diff --git a/xllm/core/common/CMakeLists.txt b/xllm/core/common/CMakeLists.txt index 445a7ea403..521e256be2 100644 --- a/xllm/core/common/CMakeLists.txt +++ b/xllm/core/common/CMakeLists.txt @@ -17,6 +17,7 @@ cc_library( types.h device_monitor.h version_singleton.h + flash_comm1_context.h SRCS etcd_client.cpp metrics.cpp @@ -24,9 +25,11 @@ cc_library( options.cpp rate_limiter.cpp device_monitor.cpp + flash_comm1_context.cpp DEPS :config util + parallel_state absl::random_random absl::strings torch diff --git a/xllm/core/common/flash_comm1_context.cpp b/xllm/core/common/flash_comm1_context.cpp new file mode 100644 index 0000000000..3ea1ec1791 --- /dev/null +++ b/xllm/core/common/flash_comm1_context.cpp @@ -0,0 +1,356 @@ +/* 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 "flash_comm1_context.h" + +#include + +#include + +#include "common/global_flags.h" +#include "framework/parallel_state/parallel_state.h" +#if defined(USE_NPU) +#include "platform/device.h" +#endif + +DEFINE_bool(enable_flashcomm1, + false, + "Enable Flash Communication 1 (FC1) sequence-parallel optimization " + "for tensor parallel inference on NPU."); + +DEFINE_int32(flashcomm1_min_prefill_tokens, + 1000, + "Minimum prefill token count to activate FC1."); + +DEFINE_int32(flashcomm1_min_decode_tokens, + 128, + "Minimum decode batch token count to activate FC1."); + +DEFINE_bool(enable_mmrs_fusion, + false, + "Enable Matmul+ReduceScatter fusion kernel for FC1."); + +DEFINE_string(mmrs_comm_mode, + "aiv", + "Communication mode for torch_npu npu_mm_reduce_scatter_base. " + "Supported values: ai_cpu, aiv, none."); + +namespace xllm { + +namespace { + +constexpr int32_t kFc1LocalTokenAlignment = 16; + +int32_t round_up_to_multiple(int32_t value, int32_t multiple) { + CHECK_GT(multiple, 0); + const int32_t remainder = value % multiple; + return remainder == 0 ? value : value + multiple - remainder; +} + +int32_t local_num_tokens_for_rank(int32_t num_tokens, + int32_t world_size, + int32_t rank) { + const int32_t base = num_tokens / world_size; + const int32_t remainder = num_tokens % world_size; + return base + (rank < remainder ? 1 : 0); +} + +int64_t shard_start_for_rank(int32_t num_tokens, + int32_t world_size, + int32_t rank) { + const int32_t base = num_tokens / world_size; + const int32_t remainder = num_tokens % world_size; + return static_cast(rank) * base + std::min(rank, remainder); +} + +torch::Tensor pad_rows_by_copy(const torch::Tensor& input, + int64_t padded_rows) { + CHECK_GE(padded_rows, input.size(0)); + if (padded_rows == input.size(0)) { + return input; + } + + auto output_shape = input.sizes().vec(); + output_shape[0] = padded_rows; + auto output = torch::empty(output_shape, input.options()); + output.slice(0, 0, input.size(0)).copy_(input); + output.slice(0, input.size(0), padded_rows).zero_(); + return output; +} + +} // namespace + +int32_t FlashComm1Context::local_num_tokens_for_rank(int32_t rank) const { + CHECK_GE(rank, 0); + CHECK_LT(rank, tp_world_size); + return xllm::local_num_tokens_for_rank( + original_num_tokens, tp_world_size, rank); +} + +std::vector FlashComm1Context::token_num_list() const { + std::vector token_nums(tp_world_size); + for (int32_t rank = 0; rank < tp_world_size; ++rank) { + token_nums[rank] = local_num_tokens_for_rank(rank); + } + return token_nums; +} + +int64_t FlashComm1Context::get_shard_start() const { + return shard_start_for_rank(original_num_tokens, tp_world_size, tp_rank); +} + +int64_t FlashComm1Context::get_shard_end() const { + return get_shard_start() + local_num_tokens; +} + +FlashComm1Context build_flash_comm1_context( + int32_t num_tokens, + bool is_prefill, + const ParallelArgs& parallel_args, + const FlashComm1Options& options) { + int32_t actual_tp_size = parallel_args.world_size() / + (parallel_args.dp_size() * parallel_args.cp_size()); + + FlashComm1Context ctx; + + if (!options.enable_flashcomm1) { + return ctx; + } + + if (!is_prefill) { + return ctx; + } + +#if defined(USE_NPU) + if (Device::type_str() != "npu") { + return ctx; + } +#else + return ctx; +#endif + + if (parallel_args.dp_size() != 1 || parallel_args.cp_size() != 1) { + return ctx; + } + + if (actual_tp_size <= 1) { + return ctx; + } + + ProcessGroup* tp_group = parallel_args.tp_group_; + if (!tp_group) { + return ctx; + } + + int32_t threshold = options.min_prefill_tokens; + + if (num_tokens <= threshold) { + return ctx; + } + + ctx.enabled = true; + ctx.tp_rank = tp_group->rank(); + ctx.tp_world_size = tp_group->world_size(); + ctx.original_num_tokens = num_tokens; + ctx.is_prefill = is_prefill; + ctx.enable_mmrs_fusion = options.enable_mmrs_fusion; + ctx.mmrs_comm_mode = options.mmrs_comm_mode; + ctx.tp_group = tp_group; + + const int32_t token_alignment = + ctx.tp_world_size * kFc1LocalTokenAlignment; + ctx.padded_num_tokens = round_up_to_multiple(num_tokens, token_alignment); + ctx.pad_size = ctx.padded_num_tokens - num_tokens; + ctx.local_num_tokens = ctx.local_num_tokens_for_rank(ctx.tp_rank); + ctx.padded_local_num_tokens = ctx.padded_num_tokens / ctx.tp_world_size; + + return ctx; +} + +FlashComm1Context build_flash_comm1_context(int32_t num_tokens, + bool is_prefill, + const ParallelArgs& parallel_args) { + FlashComm1Options options; + options.enable_flashcomm1 = FLAGS_enable_flashcomm1; + options.min_prefill_tokens = FLAGS_flashcomm1_min_prefill_tokens; + options.min_decode_tokens = FLAGS_flashcomm1_min_decode_tokens; + options.enable_mmrs_fusion = FLAGS_enable_mmrs_fusion; + options.mmrs_comm_mode = FLAGS_mmrs_comm_mode; + return build_flash_comm1_context( + num_tokens, is_prefill, parallel_args, options); +} + +torch::Tensor shard_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx) { + if (!ctx.is_sequence_sharded()) { + return input; + } + + CHECK_EQ(input.size(0), ctx.original_num_tokens); + torch::Tensor padded_input = input; + if (ctx.pad_size > 0) { + padded_input = pad_rows_by_copy(input, ctx.padded_num_tokens); + } + + const int64_t shard_start = + static_cast(ctx.tp_rank) * ctx.padded_local_num_tokens; + const int64_t shard_end = shard_start + ctx.padded_local_num_tokens; + return padded_input.slice(0, shard_start, shard_end).contiguous(); +} + +torch::Tensor gather_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx) { + if (!ctx.is_sequence_sharded()) { + return input; + } + + const int32_t current_local_size = input.size(0); + const int32_t expected_even_size = ctx.padded_num_tokens / ctx.tp_world_size; + const int32_t expected_local_size = ctx.local_num_tokens; + + std::vector token_num_list = ctx.token_num_list(); + if (ctx.pad_size > 0 && current_local_size == expected_even_size) { + for (int32_t i = 0; i < ctx.tp_world_size; ++i) { + token_num_list[i] = expected_even_size; + } + } else { + CHECK_EQ(current_local_size, expected_local_size) + << "FC1 gather expected local real-token shard size " + << expected_local_size << ", got " << current_local_size + << ", rank=" << ctx.tp_rank << ", world=" << ctx.tp_world_size; + } + + auto gathered = parallel_state::gather(input, ctx.tp_group, token_num_list); + + if (ctx.pad_size > 0 && gathered.size(0) > ctx.original_num_tokens) { + return gathered.slice(0, 0, ctx.original_num_tokens); + } + return gathered; +} + +torch::Tensor gather_and_unpad_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx) { + auto gathered = gather_sequence(input, ctx); + if (ctx.pad_size > 0) { + return gathered.slice(0, 0, ctx.original_num_tokens); + } + return gathered; +} + +torch::Tensor maybe_pad_for_reduce(const torch::Tensor& input, + const FlashComm1Context& ctx) { + int32_t current_size = input.size(0); + int32_t remainder = current_size % ctx.tp_world_size; + + if (remainder == 0) { + return input; + } + + int32_t dynamic_pad_size = ctx.tp_world_size - remainder; + return pad_rows_by_copy(input, current_size + dynamic_pad_size); +} + +namespace { + +torch::Tensor reduce_scatter_padded_local(const torch::Tensor& input, + const FlashComm1Context& ctx) { + CHECK(ctx.tp_group); + CHECK(ctx.is_sequence_sharded()); + CHECK_EQ(input.size(0), ctx.original_num_tokens) + << "FC1 row-parallel reduce_scatter expects full real-token output " + << "before communication."; + + torch::Tensor padded_input = input; + if (ctx.pad_size > 0) { + padded_input = pad_rows_by_copy(input, ctx.padded_num_tokens); + } + + auto output_shape = padded_input.sizes().vec(); + output_shape[0] = ctx.padded_local_num_tokens; + torch::Tensor output = torch::empty(output_shape, padded_input.options()); + LOG_FIRST_N(INFO, 16) + << "FC1 reduce_scatter active: rank=" << ctx.tp_rank + << ", original_tokens=" << ctx.original_num_tokens + << ", padded_tokens=" << ctx.padded_num_tokens + << ", local_tokens=" << ctx.padded_local_num_tokens; + ctx.tp_group->reduce_scatter(padded_input, output); + return output; +} + +} // namespace + +torch::Tensor maybe_pad_and_reduce(torch::Tensor input, + const FlashComm1Context& ctx, + RowParallelReduceMode mode) { + if (mode == RowParallelReduceMode::NONE) { + return input; + } + + CHECK(mode == RowParallelReduceMode::ALL_REDUCE || + mode == RowParallelReduceMode::REDUCE_SCATTER || + mode == RowParallelReduceMode::MATMUL_REDUCE_SCATTER) + << "Unsupported row-parallel reduce mode."; + + if (!ctx.is_sequence_sharded()) { + if (ctx.tp_group && ctx.tp_group->world_size() > 1) { + return parallel_state::reduce(input, ctx.tp_group); + } + return input; + } + + return reduce_scatter_padded_local(input, ctx); +} + +RowParallelReduceMode row_parallel_reduce_mode_for_fc1( + const FlashComm1Context& ctx) { + return ctx.enable_mmrs_fusion ? RowParallelReduceMode::MATMUL_REDUCE_SCATTER + : RowParallelReduceMode::REDUCE_SCATTER; +} + +torch::Tensor maybe_chunk_residual(const torch::Tensor& residual, + int32_t tp_rank, + int32_t tp_world_size) { + if (tp_world_size <= 1) { + return residual; + } + + CHECK_GE(tp_rank, 0); + CHECK_LT(tp_rank, tp_world_size); + const int32_t num_tokens = static_cast(residual.size(0)); + const int64_t start = + shard_start_for_rank(num_tokens, tp_world_size, tp_rank); + const int64_t end = + start + local_num_tokens_for_rank(num_tokens, tp_world_size, tp_rank); + return residual.slice(0, start, end).contiguous(); +} + +torch::Tensor maybe_shard_residual(const torch::Tensor& residual, + const FlashComm1Context& ctx) { + if (!ctx.is_sequence_sharded()) { + return residual; + } + if (residual.size(0) == ctx.padded_local_num_tokens) { + return residual; + } + if (residual.size(0) == ctx.original_num_tokens) { + return shard_sequence(residual, ctx); + } + CHECK_EQ(residual.size(0), ctx.padded_local_num_tokens) + << "FC1 residual layout must be either full real-token or padded local " + << "sequence shard."; + return residual; +} + +} // namespace xllm diff --git a/xllm/core/common/flash_comm1_context.h b/xllm/core/common/flash_comm1_context.h new file mode 100644 index 0000000000..73fba5aee8 --- /dev/null +++ b/xllm/core/common/flash_comm1_context.h @@ -0,0 +1,103 @@ +/* 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 "framework/parallel_state/parallel_args.h" + +namespace xllm { + +namespace layer { +struct AttentionMetadata; +} + +enum class RowParallelReduceMode : int8_t { + NONE = 0, + ALL_REDUCE = 1, + REDUCE_SCATTER = 2, + MATMUL_REDUCE_SCATTER = 3, +}; + +struct FlashComm1Context { + bool enabled = false; + int32_t tp_rank = 0; + int32_t tp_world_size = 1; + int32_t original_num_tokens = 0; + int32_t padded_num_tokens = 0; + int32_t local_num_tokens = 0; + int32_t padded_local_num_tokens = 0; + int32_t pad_size = 0; + bool is_prefill = true; + bool enable_mmrs_fusion = false; + std::string mmrs_comm_mode = "aiv"; + ProcessGroup* tp_group = nullptr; + + FlashComm1Context() = default; + + bool is_sequence_sharded() const { return enabled && tp_world_size > 1; } + + int32_t local_num_tokens_for_rank(int32_t rank) const; + std::vector token_num_list() const; + + int64_t get_shard_start() const; + int64_t get_shard_end() const; +}; + +struct FlashComm1Options { + bool enable_flashcomm1 = false; + int32_t min_prefill_tokens = 1000; + int32_t min_decode_tokens = 128; + bool enable_mmrs_fusion = false; + std::string mmrs_comm_mode = "aiv"; +}; + +FlashComm1Context build_flash_comm1_context( + int32_t num_tokens, + bool is_prefill, + const ParallelArgs& parallel_args); + +FlashComm1Context build_flash_comm1_context( + int32_t num_tokens, + bool is_prefill, + const ParallelArgs& parallel_args, + const FlashComm1Options& options); + +torch::Tensor shard_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); + +torch::Tensor gather_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); + +torch::Tensor gather_and_unpad_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); + +torch::Tensor maybe_pad_for_reduce(const torch::Tensor& input, const FlashComm1Context& ctx); + +torch::Tensor maybe_pad_and_reduce(torch::Tensor input, + const FlashComm1Context& ctx, + RowParallelReduceMode mode); + +RowParallelReduceMode row_parallel_reduce_mode_for_fc1( + const FlashComm1Context& ctx); + +torch::Tensor maybe_chunk_residual(const torch::Tensor& residual, + int32_t tp_rank, + int32_t tp_world_size); + +torch::Tensor maybe_shard_residual(const torch::Tensor& residual, + const FlashComm1Context& ctx); + +} // namespace xllm diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index a560f85c1c..a5f17ef488 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -387,3 +387,14 @@ DECLARE_bool(enable_aclnn_swiglu); DECLARE_bool(use_cpp_chat_template); DECLARE_int32(health_check_interval_ms); + +// --- Flash Communication 1 (FC1) config --- +DECLARE_bool(enable_flashcomm1); + +DECLARE_int32(flashcomm1_min_prefill_tokens); + +DECLARE_int32(flashcomm1_min_decode_tokens); + +DECLARE_bool(enable_mmrs_fusion); + +DECLARE_string(mmrs_comm_mode); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index 94500f0c60..c0edda7903 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -56,6 +56,12 @@ std::string Options::to_string() const { << ", task_type: " << task_type() << ", enable_mla: " << enable_mla() << ", enable_chunked_prefill: " << enable_chunked_prefill() << ", enable_prefill_sp: " << enable_prefill_sp() + << ", enable_flashcomm1: " << enable_flashcomm1() + << ", flashcomm1_min_prefill_tokens: " + << flashcomm1_min_prefill_tokens() + << ", flashcomm1_min_decode_tokens: " << flashcomm1_min_decode_tokens() + << ", enable_mmrs_fusion: " << enable_mmrs_fusion() + << ", mmrs_comm_mode: " << mmrs_comm_mode() << ", master_node_addr: " << master_node_addr().value_or("null") << ", instance_role: " << instance_role().to_string() << ", transfer_listen_port: " << transfer_listen_port() diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index 2b8dbca2f4..b5eb791193 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -122,6 +122,17 @@ class Options { PROPERTY(bool, enable_prefill_sp) = false; + // Flash Communication 1 (FC1) sequence-parallel optimization. + PROPERTY(bool, enable_flashcomm1) = false; + + PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 1000; + + PROPERTY(int32_t, flashcomm1_min_decode_tokens) = 128; + + PROPERTY(bool, enable_mmrs_fusion) = false; + + PROPERTY(std::string, mmrs_comm_mode) = "aiv"; + PROPERTY(std::optional, master_node_addr); PROPERTY(int32_t, nnodes) = 1; diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 17b8d55685..89737184dc 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -240,6 +240,12 @@ Master::Master(const Options& options, EngineType type) .task_type(options.task_type()) .enable_mla(options_.enable_mla()) .enable_prefill_sp(options_.enable_prefill_sp()) + .enable_flashcomm1(options_.enable_flashcomm1()) + .flashcomm1_min_prefill_tokens( + options_.flashcomm1_min_prefill_tokens()) + .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) + .enable_mmrs_fusion(options_.enable_mmrs_fusion()) + .mmrs_comm_mode(options_.mmrs_comm_mode()) .npu_kernel_backend(options_.npu_kernel_backend()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .enable_offline_inference(options_.enable_offline_inference()) @@ -313,6 +319,12 @@ Master::Master(const Options& options, EngineType type) .dp_size(options.dp_size()) .ep_size(options.ep_size()) .enable_prefill_sp(options_.enable_prefill_sp()) + .enable_flashcomm1(options_.enable_flashcomm1()) + .flashcomm1_min_prefill_tokens( + options_.flashcomm1_min_prefill_tokens()) + .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) + .enable_mmrs_fusion(options_.enable_mmrs_fusion()) + .mmrs_comm_mode(options_.mmrs_comm_mode()) .cp_size(options.cp_size()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .max_tokens_per_batch(options_.max_tokens_per_batch()) @@ -368,6 +380,12 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .enable_prefill_sp(options_.enable_prefill_sp()) + .enable_flashcomm1(options_.enable_flashcomm1()) + .flashcomm1_min_prefill_tokens( + options_.flashcomm1_min_prefill_tokens()) + .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) + .enable_mmrs_fusion(options_.enable_mmrs_fusion()) + .mmrs_comm_mode(options_.mmrs_comm_mode()) .cp_size(options_.cp_size()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .max_tokens_per_batch(options_.max_tokens_per_batch()) diff --git a/xllm/core/framework/model_context.cpp b/xllm/core/framework/model_context.cpp index 822c44faa5..b60104ed13 100644 --- a/xllm/core/framework/model_context.cpp +++ b/xllm/core/framework/model_context.cpp @@ -96,6 +96,7 @@ ModelContext ModelContext::with_parallel_args( #endif derived.model_id_ = model_id_; derived.optimization_config_ = optimization_config_; + derived.flash_comm1_options_ = flash_comm1_options_; return derived; } diff --git a/xllm/core/framework/model_context.h b/xllm/core/framework/model_context.h index 3a321a887e..c466de5dbe 100644 --- a/xllm/core/framework/model_context.h +++ b/xllm/core/framework/model_context.h @@ -24,6 +24,7 @@ limitations under the License. #include #include +#include "core/common/flash_comm1_context.h" #include "core/framework/model/model_args.h" #include "core/framework/parallel_state/parallel_args.h" #include "core/framework/quant_args.h" @@ -82,6 +83,14 @@ class ModelContext { return optimization_config_; } + const FlashComm1Options& get_flash_comm1_options() const { + return flash_comm1_options_; + } + + void set_flash_comm1_options(const FlashComm1Options& options) { + flash_comm1_options_ = options; + } + ModelContext with_parallel_args(const ParallelArgs& parallel_args) const; #if defined(USE_NPU) @@ -109,6 +118,7 @@ class ModelContext { ParallelArgs parallel_args_; torch::TensorOptions tensor_options_; OptimizationConfig optimization_config_; + FlashComm1Options flash_comm1_options_; #if defined(USE_NPU) // used for npu atb diff --git a/xllm/core/kernels/npu/CMakeLists.txt b/xllm/core/kernels/npu/CMakeLists.txt index 76327e9bb8..413902fcfb 100644 --- a/xllm/core/kernels/npu/CMakeLists.txt +++ b/xllm/core/kernels/npu/CMakeLists.txt @@ -15,6 +15,7 @@ cc_library( attention.cpp fused_layernorm.cpp matmul.cpp + matmul_reduce_scatter.cpp npu_causal_conv1d.cpp npu_fused_infer_attention.cpp npu_gemma_rms_norm.cpp diff --git a/xllm/core/kernels/npu/matmul.cpp b/xllm/core/kernels/npu/matmul.cpp index 71d3408959..eee3d9f13e 100644 --- a/xllm/core/kernels/npu/matmul.cpp +++ b/xllm/core/kernels/npu/matmul.cpp @@ -14,7 +14,6 @@ limitations under the License. ==============================================================================*/ #include "npu_ops_api.h" -#include "ops_npu/npu_ops.h" namespace xllm::kernel::npu { diff --git a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp new file mode 100644 index 0000000000..af686290c9 --- /dev/null +++ b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp @@ -0,0 +1,190 @@ +/* Copyright 2025-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 "npu_ops_api.h" + +#include +#include + +#include + +#include "core/framework/parallel_state/process_group.h" + +namespace xllm::kernel::npu { +namespace { + +std::string matmul_reduce_scatter_reject_reason( + const torch::Tensor& a, + const torch::Tensor& b, + const std::optional& bias, + const std::optional& output, + ProcessGroup* process_group) { + if (process_group == nullptr) { + return "process_group is null"; + } + if (!output.has_value() || !output->defined()) { + return "output tensor is missing"; + } + if (a.dim() != 2 || b.dim() != 2 || output->dim() != 2) { + std::ostringstream oss; + oss << "expected 2D tensors, got a_dim=" << a.dim() + << ", b_dim=" << b.dim() << ", output_dim=" << output->dim(); + return oss.str(); + } + if (a.scalar_type() != at::kHalf && a.scalar_type() != at::kBFloat16) { + std::ostringstream oss; + oss << "unsupported input dtype=" << a.scalar_type(); + return oss.str(); + } + if (a.scalar_type() != b.scalar_type() || + a.scalar_type() != output->scalar_type()) { + std::ostringstream oss; + oss << "dtype mismatch: a=" << a.scalar_type() + << ", b=" << b.scalar_type() + << ", output=" << output->scalar_type(); + return oss.str(); + } + if (bias.has_value() && bias->defined() && + bias->scalar_type() != a.scalar_type()) { + std::ostringstream oss; + oss << "bias dtype mismatch: bias=" << bias->scalar_type() + << ", input=" << a.scalar_type(); + return oss.str(); + } + if (a.size(1) != b.size(0)) { + std::ostringstream oss; + oss << "matmul K mismatch: a=" << a.sizes() << ", b=" << b.sizes(); + return oss.str(); + } + if (output->size(1) != b.size(1)) { + std::ostringstream oss; + oss << "output N mismatch: output=" << output->sizes() + << ", b=" << b.sizes(); + return oss.str(); + } + return ""; +} + +torch::Tensor matmul_kn(const torch::Tensor& a, + const torch::Tensor& b, + const std::optional& bias) { + torch::Tensor out = torch::matmul(a, b); + if (bias.has_value() && bias->defined()) { + out = out + bias.value(); + } + return out; +} + +bool should_use_ai_cpu_for_aiv_risky_shape(const torch::Tensor& a, + const torch::Tensor& b, + ProcessGroup* process_group) { + if (process_group == nullptr || a.scalar_type() != at::kBFloat16 || + a.dim() != 2 || b.dim() != 2 || a.size(0) != 2048 || + b.size(1) != 5120) { + return false; + } + const int64_t world_size = process_group->world_size(); + const int64_t k = a.size(1); + if (world_size == 2) { + return k == 3072 || k == 8704; + } + if (world_size == 4) { + return k == 4352; + } + if (world_size == 8) { + return k == 2176; + } + return false; +} + +} // namespace + +torch::Tensor matmul_reduce_scatter( + const torch::Tensor& a, + const torch::Tensor& b, + const std::optional& bias, + const std::optional& output, + ProcessGroup* process_group, + const std::string& reduce_op, + int64_t comm_turn, + int64_t stream_mode, + const std::string& comm_mode) { + const std::string reject_reason = + matmul_reduce_scatter_reject_reason(a, b, bias, output, process_group); + if (!reject_reason.empty()) { + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: " << reject_reason + << "; fallback to matmul + reduce_scatter path. a=" << a.sizes() + << ", b=" << b.sizes() + << ", output=" + << (output.has_value() && output->defined() ? output->sizes() + : c10::IntArrayRef{}); + return matmul_kn(a, b, bias); + } + + std::string group = process_group->hccl_comm_name(/*init_comm=*/true); + if (group.empty()) { + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: HCCL group name is empty; fallback to " + "matmul + reduce_scatter path."; + return matmul_kn(a, b, bias); + } + + std::string effective_comm_mode = comm_mode; + if ((comm_mode.empty() || comm_mode == "aiv") && + should_use_ai_cpu_for_aiv_risky_shape(a, b, process_group)) { + effective_comm_mode = "ai_cpu"; + LOG_FIRST_N(WARNING, 16) + << "FC1 MMRS comm_mode switched from aiv to ai_cpu for known risky " + "torch_npu shape: a=" + << a.sizes() << ", b=" << b.sizes() + << ", world_size=" << process_group->world_size() + << ", dtype=" << a.scalar_type(); + } + + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS torch_npu hit: a=" << a.sizes() << ", b=" << b.sizes() + << ", expected_out=" << output->sizes() << ", dtype=" << a.scalar_type() + << ", rank=" << process_group->rank() + << ", world_size=" << process_group->world_size() + << ", reduce_op=" << (reduce_op.empty() ? "sum" : reduce_op) + << ", comm_turn=" << comm_turn << ", stream_mode=" << stream_mode + << ", comm_mode=" << (comm_mode.empty() ? "none" : comm_mode) + << ", effective_comm_mode=" + << (effective_comm_mode.empty() ? "none" : effective_comm_mode); + + std::optional torch_comm_mode = std::nullopt; + if (effective_comm_mode == "ai_cpu" || effective_comm_mode == "aiv") { + torch_comm_mode = c10::string_view(effective_comm_mode); + } else if (!effective_comm_mode.empty() && effective_comm_mode != "none") { + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu unsupported comm_mode=" << effective_comm_mode + << "; using torch_npu default comm_mode."; + } + return at_npu::native::custom_ops::npu_mm_reduce_scatter_base( + a, + b, + group, + process_group->world_size(), + reduce_op.empty() ? c10::string_view("sum") : c10::string_view(reduce_op), + bias, + /*x1_scale=*/std::nullopt, + /*x2_scale=*/std::nullopt, + comm_turn, + /*output_dtype=*/std::nullopt, + torch_comm_mode); +} + +} // 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 762f4875ec..aaf936d44e 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -23,6 +23,10 @@ limitations under the License. #include "custom_functions_npu/atb_common.h" +namespace xllm { +class ProcessGroup; +} // namespace xllm + namespace xllm::kernel::npu { void reshape_paged_cache(torch::Tensor& key, @@ -89,6 +93,17 @@ torch::Tensor matmul(const torch::Tensor& a, const torch::Tensor& b, const std::optional& bias); +torch::Tensor matmul_reduce_scatter( + const torch::Tensor& a, + const torch::Tensor& b, + const std::optional& bias, + const std::optional& output, + ProcessGroup* process_group, + const std::string& reduce_op, + int64_t comm_turn, + int64_t stream_mode, + const std::string& comm_mode); + torch::Tensor active(const torch::Tensor& input, const std::string& act_mode); torch::Tensor rms_norm(const torch::Tensor& input, diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index a7a915c52b..d823bb9405 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -413,6 +413,26 @@ torch::Tensor matmul(MatmulParams& params) { #endif } +torch::Tensor matmul_reduce_scatter(MatmulReduceScatterParams& params) { +#if defined(USE_NPU) + return npu::matmul_reduce_scatter(params.a, + params.b, + params.bias, + params.output, + params.process_group, + params.reduce_op, + params.comm_turn, + params.stream_mode, + params.comm_mode); +#else + MatmulParams matmul_params; + matmul_params.a = params.a; + matmul_params.b = params.b; + matmul_params.bias = params.bias; + return matmul(matmul_params); +#endif +} + torch::Tensor quant_matmul(QuantMatmulParams& params) { #if defined(USE_NPU) return npu::quant_matmul(params.x1, diff --git a/xllm/core/kernels/ops_api.h b/xllm/core/kernels/ops_api.h index 1ead17c396..d5202c0acc 100644 --- a/xllm/core/kernels/ops_api.h +++ b/xllm/core/kernels/ops_api.h @@ -47,6 +47,8 @@ std::tuple rms_norm_dynamic_quant( torch::Tensor matmul(MatmulParams& params); +torch::Tensor matmul_reduce_scatter(MatmulReduceScatterParams& params); + torch::Tensor quant_matmul(QuantMatmulParams& params); torch::Tensor quantize(NpuQuantizeParams& params); diff --git a/xllm/core/kernels/param.h b/xllm/core/kernels/param.h index 79f2253d71..028e84cd5b 100644 --- a/xllm/core/kernels/param.h +++ b/xllm/core/kernels/param.h @@ -21,6 +21,10 @@ limitations under the License. #include #include +namespace xllm { +class ProcessGroup; +} // namespace xllm + namespace xllm::layer { struct AttentionMetadata; } // namespace xllm::layer @@ -312,6 +316,26 @@ struct MatmulParams { double beta = 0.0; }; +struct MatmulReduceScatterParams { + // RFC-visible inputs. For the currently wired BF16/FP16 path, b is passed in + // K-N layout to match torch_npu::npu_mm_reduce_scatter_base. + torch::Tensor a; + torch::Tensor b; + std::optional bias; + ProcessGroup* process_group = nullptr; + int64_t original_num_tokens = 0; + std::optional deq_scale; + std::optional output_dtype; + + // Optional compatibility fields used by callers for expected shape checks. + // The torch_npu op allocates and returns the output tensor itself. + std::optional output; + std::string reduce_op = "sum"; + int64_t comm_turn = 0; + int64_t stream_mode = 1; + std::string comm_mode = "ai_cpu"; +}; + // Quantized matmul parameters (NPU aclnnQuantMatmulV4 path). struct QuantMatmulParams { // Quantized activation tensor. Typical dtype: int8. diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index 78ed058521..cd2232aa9d 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -17,9 +17,9 @@ limitations under the License. #include +#include "common/flash_comm1_context.h" #include "kernels/ops_api.h" #include "platform/device.h" -#include "platform/platform.h" namespace xllm { namespace layer { @@ -94,25 +94,44 @@ DenseMLPImpl::DenseMLPImpl(int64_t hidden_size, down_proj_extra_args)); } -torch::Tensor DenseMLPImpl::forward(const torch::Tensor& hidden_states) { - // input shape: [num_tokens, hidden_size] - auto gate_up = gate_up_proj_->forward(hidden_states); +torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states, + const FlashComm1Context* fc1_ctx) { + torch::Tensor h = hidden_states; + + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + h = gather_sequence(hidden_states, *fc1_ctx); + } + + auto gate_up = gate_up_proj_->forward(h); if (is_smoothquant_) { - // For w8a8 quantization, the active operation is fused with the down_proj - return down_proj_->forward(gate_up); - } else { - torch::Tensor output; - if (!Platform::is_npu()) { - int64_t batch_size = gate_up.sizes()[0]; - output = torch::empty( - {batch_size, intermediate_size_ / process_group_->world_size()}, - gate_up.options()); + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS callsite DenseMLP.down_proj(smoothquant): input=" + << gate_up.sizes(); + return down_proj_->forward( + gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); } + return down_proj_->forward(gate_up); + } + + torch::Tensor output; + if (Device::type_str() != "npu") { + int64_t batch_size = gate_up.sizes()[0]; + output = torch::empty( + {batch_size, intermediate_size_ / process_group_->world_size()}, + gate_up.options()); + } + + act_->forward(gate_up, output); - act_->forward(gate_up, output); - return down_proj_->forward(output); + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS callsite DenseMLP.down_proj: input=" << output.sizes(); + return down_proj_->forward( + output, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); } + return down_proj_->forward(output); } void DenseMLPImpl::load_state_dict(const StateDict& state_dict) { diff --git a/xllm/core/layers/common/dense_mlp.h b/xllm/core/layers/common/dense_mlp.h index a6f6222ae4..d2cdeff310 100644 --- a/xllm/core/layers/common/dense_mlp.h +++ b/xllm/core/layers/common/dense_mlp.h @@ -20,6 +20,7 @@ limitations under the License. #include #include "activation.h" +#include "common/flash_comm1_context.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" #include "framework/quant_args.h" @@ -44,7 +45,8 @@ class DenseMLPImpl : public torch::nn::Module { const std::string& module_prefix = "", double swiglu_limit = 0.0); - torch::Tensor forward(const torch::Tensor& hidden_states); + torch::Tensor forward(torch::Tensor hidden_states, + const FlashComm1Context* fc1_ctx = nullptr); void load_state_dict(const StateDict& state_dict); void load_state_dict(const StateDict& state_dict, diff --git a/xllm/core/layers/common/linear.cpp b/xllm/core/layers/common/linear.cpp index 728ff16e71..50bcc1ebc3 100644 --- a/xllm/core/layers/common/linear.cpp +++ b/xllm/core/layers/common/linear.cpp @@ -61,6 +61,21 @@ inline bool is_unfused_checkpoint(const std::vector& scales) { scales.back() > std::numeric_limits::lowest(); } +torch::Tensor pad_rows_by_copy(const torch::Tensor& input, + int64_t padded_rows) { + CHECK_GE(padded_rows, input.size(0)); + if (padded_rows == input.size(0)) { + return input; + } + + auto output_shape = input.sizes().vec(); + output_shape[0] = padded_rows; + auto output = torch::empty(output_shape, input.options()); + output.slice(0, 0, input.size(0)).copy_(input); + output.slice(0, input.size(0), padded_rows).zero_(); + return output; +} + // Realigns FP8 partitions to a unified global scale to enable fusion. // Logic: // 1. Recover original values (FP8 -> FP16) using partition-specific scales. @@ -245,6 +260,27 @@ bool is_w8a8_quant( resolved_weight_quant_method.value() == "w8a8"; } +bool wants_mmrs(RowParallelReduceMode reduce_mode) { + return reduce_mode == RowParallelReduceMode::MATMUL_REDUCE_SCATTER; +} + +void log_mmrs_quant_skip(RowParallelReduceMode reduce_mode, + const FlashComm1Context* fc1_ctx, + const char* quant_path, + const torch::Tensor& input) { + if (!wants_mmrs(reduce_mode)) { + return; + } + LOG_FIRST_N(WARNING, 16) + << "FC1 MMRS skipped in row-parallel " << quant_path + << " path: fused matmul_reduce_scatter is currently wired only for " + "non-quant linear. input=" << input.sizes() + << ", sequence_sharded=" + << (fc1_ctx != nullptr && fc1_ctx->is_sequence_sharded()) + << ", enable_mmrs_fusion=" + << (fc1_ctx != nullptr && fc1_ctx->enable_mmrs_fusion); +} + torch::Dtype get_w8a8_deq_scale_dtype(const torch::TensorOptions& options) { const torch::Dtype dtype = c10::typeMetaToScalarType(options.dtype()); if (dtype == torch::kFloat16) { @@ -1521,11 +1557,238 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input) { return output; } +torch::Tensor RowParallelLinearImpl::mmrs_weight_transposed() const { + CHECK(weight_.defined()) << "weight is required for MMRS."; + const bool valid = + mmrs_weight_t_.defined() && mmrs_weight_t_.device() == weight_.device() && + mmrs_weight_t_.scalar_type() == weight_.scalar_type() && + mmrs_weight_t_.size(0) == weight_.size(1) && + mmrs_weight_t_.size(1) == weight_.size(0); + if (!valid) { + mmrs_weight_t_ = weight_.transpose(0, 1).contiguous(); + } + return mmrs_weight_t_; +} + +torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, + RowParallelReduceMode reduce_mode, + const FlashComm1Context* fc1_ctx) { + auto bias = bias_.defined() && rank_ == 0 + ? std::optional(bias_) + : std::nullopt; + + const bool skip_scatter = fc1_ctx && fc1_ctx->is_sequence_sharded(); + + torch::Tensor output; + if (quant_args_.quant_method() == kQuantMethodSmoothquant) { + log_mmrs_quant_skip(reduce_mode, fc1_ctx, "smoothquant", input); + CHECK(smooth_.defined()) << "smooth is required for smoothquant."; + CHECK(qweight_.defined()) << "qweight is required for smoothquant."; + CHECK(per_channel_scale_.defined()) + << "per_channel_scale is required for smoothquant."; + + torch::Tensor quantized_input; + torch::Tensor input_scale; + + if (!input_is_parallelized_ && !skip_scatter) { + input = xllm::parallel_state::scatter(input, process_group_); + } + + xllm::kernel::ScaledQuantizeParams quantize_params; + quantize_params.x = input; + quantize_params.smooth = smooth_; + quantize_params.zero = std::nullopt; + quantize_params.token_count = std::nullopt; + quantize_params.gather_index = std::nullopt; + quantize_params.gather_index_start_position = std::nullopt; + quantize_params.output = std::nullopt; + quantize_params.output_scale = std::nullopt; + quantize_params.act_mode = linear_extra_args_.act_mode; + quantize_params.active_coef = 1.0; + quantize_params.is_gated = linear_extra_args_.is_gated; + + std::tie(quantized_input, input_scale) = + xllm::kernel::scaled_quantize(quantize_params); + + xllm::kernel::ScaledMatmulParams matmul_params; + matmul_params.a = quantized_input; + matmul_params.b = qweight_; + matmul_params.a_scale = input_scale; + matmul_params.b_scale = per_channel_scale_; + matmul_params.output_dtype = output_dtype_; + matmul_params.bias = bias; + matmul_params.c = std::nullopt; + matmul_params.act_mode = "none"; + matmul_params.quant_bit_size = 8; + matmul_params.alpha = 1.0; + matmul_params.beta = 0.0; + matmul_params.use_hp_active = false; + matmul_params.a_quant_bit_size = 8; + matmul_params.a_calib = std::nullopt; + matmul_params.b_calib = std::nullopt; + matmul_params.output = std::nullopt; + + output = xllm::kernel::scaled_matmul(matmul_params); + } else if (quant_args_.quant_method() == kQuantMethodFp8) { + log_mmrs_quant_skip(reduce_mode, fc1_ctx, "fp8", input); + CHECK(!quant_args_.activation_dynamic()) + << "FP8 quantization does not support activation_dynamic yet"; + + if (!input_is_parallelized_ && !skip_scatter) { + input = xllm::parallel_state::scatter(input, process_group_); + } + + auto scale = input_scale_.defined() + ? std::optional(input_scale_) + : std::nullopt; + output = fp8_linear_forward( + input, weight_, weight_scale_, scale, bias, output_dtype_); + } else if (is_w8a8_quant(resolved_weight_quant_method_)) { + log_mmrs_quant_skip(reduce_mode, fc1_ctx, "w8a8", input); + CHECK(input_scale_is_loaded_ && input_scale_.defined()) + << "input_scale is required for w8a8 quant matmul."; + CHECK(input_offset_is_loaded_ && input_offset_.defined()) + << "input_offset is required for w8a8 quant matmul."; + CHECK(deq_scale_is_loaded_ && deq_scale_.defined()) + << "deq_scale is required for w8a8 quant matmul."; + if (!input_is_parallelized_ && !skip_scatter) { + input = xllm::parallel_state::scatter(input, process_group_); + } + auto quant_bias = quant_bias_is_loaded_ && quant_bias_.defined() + ? std::optional(quant_bias_) + : std::nullopt; + output = npu_w8a8_linear_forward(input, + weight_, + input_scale_, + input_offset_, + deq_scale_, + quant_bias, + output_dtype_); + } else if (is_w8a8_dynamic_quant(resolved_weight_quant_method_)) { + log_mmrs_quant_skip(reduce_mode, fc1_ctx, "w8a8_dynamic", input); + if (!input_is_parallelized_ && !skip_scatter) { + input = xllm::parallel_state::scatter(input, process_group_); + } + auto weight_scale = weight_scale_is_loaded_ + ? std::optional(weight_scale_) + : std::nullopt; + CHECK(weight_scale.has_value() && weight_scale.value().defined()) + << "weight_scale is required for w8a8_dynamic quant matmul."; +#if defined(USE_DCU) + output = dcu_w8a8_dynamic_linear_forward( + input, weight_, weight_scale.value(), bias, output_dtype_); +#elif defined(USE_NPU) + output = npu_w8a8_dynamic_linear_forward( + input, weight_, weight_scale.value(), bias, output_dtype_); +#endif + } else { + if (!input_is_parallelized_ && !skip_scatter) { + input = xllm::parallel_state::scatter(input, process_group_); + } + if (wants_mmrs(reduce_mode) && fc1_ctx && fc1_ctx->is_sequence_sharded() && + fc1_ctx->enable_mmrs_fusion) { + bool can_try_mmrs = input.dim() == 2 && + input.size(0) == fc1_ctx->original_num_tokens && + (!bias.has_value() || fc1_ctx->pad_size == 0); + if (can_try_mmrs) { + torch::Tensor mmrs_input = input; + if (fc1_ctx->pad_size > 0) { + mmrs_input = pad_rows_by_copy(input, fc1_ctx->padded_num_tokens); + } + + auto output_shape = mmrs_input.sizes().vec(); + output_shape[0] = fc1_ctx->padded_local_num_tokens; + output_shape[1] = weight_.size(0); + torch::Tensor mmrs_output = torch::empty(output_shape, input.options()); + + xllm::kernel::MatmulReduceScatterParams mmrs_params; + mmrs_params.a = mmrs_input; + mmrs_params.b = mmrs_weight_transposed(); + mmrs_params.bias = bias; + mmrs_params.process_group = process_group_; + mmrs_params.original_num_tokens = fc1_ctx->original_num_tokens; + mmrs_params.output = mmrs_output; + mmrs_params.comm_mode = fc1_ctx->mmrs_comm_mode; + output = xllm::kernel::matmul_reduce_scatter(mmrs_params); + if (output.sizes() == torch::IntArrayRef(output_shape)) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS row-parallel fused output accepted: input=" + << input.sizes() << ", weight=" << weight_.sizes() + << ", output=" << output.sizes(); + return output; + } + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS returned non-local shape; fallback reduction will run. " + << "input=" << input.sizes() << ", weight=" << weight_.sizes() + << ", returned_output=" << output.sizes() + << ", expected_local_output=" << output_shape; + if (fc1_ctx->pad_size > 0 && + output.size(0) == fc1_ctx->padded_num_tokens) { + output = output.slice(0, 0, fc1_ctx->original_num_tokens) + .contiguous(); + } + } else { + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS skipped for unsupported row-parallel shape; fallback " + "to matmul + reduce_scatter. input=" << input.sizes() + << ", weight=" << weight_.sizes() + << ", original_num_tokens=" << fc1_ctx->original_num_tokens + << ", pad_size=" << fc1_ctx->pad_size + << ", has_bias=" << bias.has_value() + << ", input_dim=" << input.dim(); + } + + if (!output.defined()) { + xllm::kernel::MatmulParams matmul_params; + matmul_params.a = input; + matmul_params.b = weight_; + matmul_params.bias = bias; + output = xllm::kernel::matmul(matmul_params); + } + } else { + if (wants_mmrs(reduce_mode)) { + LOG_FIRST_N(WARNING, 16) + << "FC1 MMRS skipped before row-parallel matmul: fc1_ctx=" + << (fc1_ctx != nullptr) + << ", sequence_sharded=" + << (fc1_ctx != nullptr && fc1_ctx->is_sequence_sharded()) + << ", enable_mmrs_fusion=" + << (fc1_ctx != nullptr && fc1_ctx->enable_mmrs_fusion) + << ", reduce_mode=" << static_cast(reduce_mode) + << ", input=" << input.sizes(); + } + xllm::kernel::MatmulParams matmul_params; + matmul_params.a = input; + matmul_params.b = weight_; + matmul_params.bias = bias; + output = xllm::kernel::matmul(matmul_params); + } + } + + if (reduce_mode == RowParallelReduceMode::NONE) { + return output; + } + + if ((reduce_mode == RowParallelReduceMode::REDUCE_SCATTER || + reduce_mode == RowParallelReduceMode::MATMUL_REDUCE_SCATTER) && + fc1_ctx) { + FlashComm1Context ctx_copy = *fc1_ctx; + ctx_copy.tp_group = process_group_; + return maybe_pad_and_reduce(output, ctx_copy, reduce_mode); + } + + if (enable_result_reduction_ && world_size_ > 1) { + output = xllm::parallel_state::reduce(output, process_group_); + } + return output; +} + // load the weight from the checkpoint void RowParallelLinearImpl::load_state_dict(const StateDict& state_dict) { if (state_dict.size() == 0) { return; } + mmrs_weight_t_ = torch::Tensor(); const int64_t rank = world_size_ == 1 ? 0 : rank_; const int64_t world_size = world_size_; resolve_weight_quant_method_for_linear_load( @@ -1589,6 +1852,7 @@ void RowParallelLinearImpl::load_state_dict(const StateDict& state_dict) { if (bias_.defined()) { LOAD_WEIGHT(bias); } + mmrs_weight_t_ = torch::Tensor(); } // Linear layer with row parallelism. diff --git a/xllm/core/layers/common/linear.h b/xllm/core/layers/common/linear.h index ad3186be2f..35db282059 100644 --- a/xllm/core/layers/common/linear.h +++ b/xllm/core/layers/common/linear.h @@ -19,6 +19,7 @@ limitations under the License. #include #include "core/framework/model_context.h" +#include "common/flash_comm1_context.h" #include "framework/parallel_state/parallel_args.h" #include "framework/quant_args.h" #include "framework/state_dict/state_dict.h" @@ -259,6 +260,10 @@ class RowParallelLinearImpl : public torch::nn::Module { torch::Tensor forward(torch::Tensor input); + torch::Tensor forward(torch::Tensor input, + RowParallelReduceMode reduce_mode, + const FlashComm1Context* fc1_ctx); + // load the weight from the checkpoint void load_state_dict(const StateDict& state_dict); @@ -283,6 +288,8 @@ class RowParallelLinearImpl : public torch::nn::Module { ProcessGroup* process_group() const { return process_group_; } private: + torch::Tensor mmrs_weight_transposed() const; + // parameter members, must be registered // we allocate the transpose since linear performs XA^T. // A^T: [out_features, in_features_per_partition] @@ -325,6 +332,7 @@ class RowParallelLinearImpl : public torch::nn::Module { at::ScalarType output_dtype_; LinearExtraArgs linear_extra_args_; std::optional resolved_weight_quant_method_; + mutable torch::Tensor mmrs_weight_t_; }; TORCH_MODULE(RowParallelLinear); diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index 7312550dda..06aea7069b 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -533,10 +533,23 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params) { + return forward(hidden_states, attn_metadata, kv_cache, input_params, nullptr); +} + +torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const FlashComm1Context* fc1_ctx) { + torch::Tensor h = hidden_states; + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + h = gather_sequence(hidden_states, *fc1_ctx); + } + // Save original hidden_states size for potential padding later - const int64_t original_num_tokens = hidden_states.size(0); - auto [qkvz_padded, ba_padded] = - project_padded_inputs(hidden_states, attn_metadata); + const int64_t original_num_tokens = h.size(0); + auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata); int64_t batch_size = qkvz_padded.size(0); int64_t seq_len = qkvz_padded.size(1); @@ -909,6 +922,13 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( rearranged_norm = rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); } + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS callsite Qwen3GatedDeltaNet.o_proj: input=" + << rearranged_norm.sizes(); + return o_proj_->forward( + rearranged_norm, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + } return o_proj_->forward(rearranged_norm); } diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h index afbda14b77..a8350356d7 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h @@ -22,6 +22,7 @@ limitations under the License. #include #include "attention.h" +#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" @@ -50,6 +51,12 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { KVCache& kv_cache, const ModelInputParams& input_params); + torch::Tensor forward(const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const FlashComm1Context* fc1_ctx); + protected: virtual std::pair project_decode_inputs( const torch::Tensor& hidden_states) = 0; diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index e5838f13cf..6134cb1d86 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -20,6 +20,8 @@ limitations under the License. #include #include +#include "common/flash_comm1_context.h" + namespace xllm { namespace layer { @@ -145,7 +147,23 @@ torch::Tensor Qwen3NextAttentionImpl::forward( const AttentionMetadata& attn_metadata, KVCache& kv_cache, const torch::Tensor& mrope_cos_sin) { - auto qkv = qkv_proj_->forward(hidden_states); + return forward(positions, hidden_states, attn_metadata, kv_cache, mrope_cos_sin, nullptr); +} + +torch::Tensor Qwen3NextAttentionImpl::forward( + const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin, + const FlashComm1Context* fc1_ctx) { + torch::Tensor h = hidden_states; + + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + h = gather_sequence(hidden_states, *fc1_ctx); + } + + auto qkv = qkv_proj_->forward(h); if (use_fused_qkv_) { const int64_t T = qkv.size(0); @@ -169,10 +187,17 @@ torch::Tensor Qwen3NextAttentionImpl::forward( auto out = std::get<0>( attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache)); out = out * torch::sigmoid(gate.view({T, q_size_})); + + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS callsite Qwen3NextAttention.o_proj(mrope): input=" + << out.sizes(); + return o_proj_->forward( + out, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + } return o_proj_->forward(out); } - // Fallback path: weight-reordered layout [Q | G | K | V] torch::Tensor q, k, v; torch::Tensor gate; @@ -199,6 +224,14 @@ torch::Tensor Qwen3NextAttentionImpl::forward( if (attn_output_gate_) { out = out * torch::sigmoid(gate); } + + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + LOG_FIRST_N(INFO, 16) + << "FC1 MMRS callsite Qwen3NextAttention.o_proj: input=" + << out.sizes(); + return o_proj_->forward( + out, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + } return o_proj_->forward(out); } diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.h b/xllm/core/layers/npu_torch/qwen3_next_attention.h index 45347fb919..dbfd34e37c 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.h +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.h @@ -20,6 +20,7 @@ limitations under the License. #include #include "attention.h" +#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" @@ -42,11 +43,18 @@ class Qwen3NextAttentionImpl : public torch::nn::Module { const torch::TensorOptions& options, int32_t layer_id); +torch::Tensor forward(const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin); + torch::Tensor forward(const torch::Tensor& positions, - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const torch::Tensor& mrope_cos_sin); + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin, + const FlashComm1Context* fc1_ctx); torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const; diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp index 56280be4ae..b3a0ea4ecd 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp @@ -19,6 +19,8 @@ limitations under the License. #include #include +#include "common/flash_comm1_context.h" + namespace xllm { namespace layer { @@ -114,30 +116,64 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( KVCache& kv_cache, const ModelInputParams& input_params, const torch::Tensor& mrope_cos_sin) { - // Pre-attention norm + return forward(x, + residual, + positions, + attn_metadata, + kv_cache, + input_params, + mrope_cos_sin, + nullptr); +} + +torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( + torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin, + const FlashComm1Context* fc1_ctx) { if (!residual.has_value()) { residual = x; x = std::get<0>(input_norm_->forward(x)); } else { + if (fc1_ctx && fc1_ctx->is_sequence_sharded() && + residual.value().size(0) != x.size(0)) { + residual = maybe_shard_residual(residual.value(), *fc1_ctx); + } + if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + CHECK_EQ(residual.value().size(0), x.size(0)) + << "FC1 input residual and hidden states must share the same " + << "padded local sequence layout."; + } std::tie(x, residual) = input_norm_->forward(x, residual); } - // Attention if (attention_) { x = attention_->forward( - positions, x, attn_metadata, kv_cache, mrope_cos_sin); + positions, x, attn_metadata, kv_cache, mrope_cos_sin, fc1_ctx); } else { - x = linear_attention_->forward(x, attn_metadata, kv_cache, input_params); + x = linear_attention_->forward( + x, attn_metadata, kv_cache, input_params, fc1_ctx); + } + + // Before post_norm, ensure residual shape matches x shape + if (fc1_ctx && fc1_ctx->is_sequence_sharded() && residual.has_value() && + residual.value().size(0) != x.size(0)) { + residual = maybe_shard_residual(residual.value(), *fc1_ctx); + CHECK_EQ(residual.value().size(0), x.size(0)) + << "FC1 post-attention residual and hidden states must share the same " + << "padded local sequence layout."; } - // Post-attention norm std::tie(x, residual) = post_norm_->forward(x, residual); - // MLP forward if (moe_mlp_) { x = moe_mlp_(x, input_params); } else { - x = mlp_(x); + x = mlp_->forward(x, fc1_ctx); } return x; diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h index fb6d3a6c83..f50b4f1b2f 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h @@ -20,6 +20,7 @@ limitations under the License. #include #include +#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" #include "framework/model_context.h" @@ -37,13 +38,21 @@ class Qwen3HybridDecoderLayerModule : public torch::nn::Module { public: virtual void load_state_dict(const StateDict& state_dict) = 0; virtual void verify_loaded_weights(const std::string& prefix) const = 0; +virtual torch::Tensor forward(torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) = 0; virtual torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) = 0; + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin, + const FlashComm1Context* fc1_ctx) = 0; virtual torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const { return {}; @@ -64,13 +73,22 @@ class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule { void verify_loaded_weights(const std::string& prefix) const override; +torch::Tensor forward(torch::Tensor& x, + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) override; + torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) override; + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin, + const FlashComm1Context* fc1_ctx) override; torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const override; diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 96c6b4b2d0..c567317fe7 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -250,9 +250,15 @@ std::optional LLMWorkerImpl::step_internal( eplb_executor_->eplb_execute(input.input_params.expert.eplb_info); } + ModelInputParams model_input_params = input.input_params; + if (sampling_params.selected_token_idxes.defined()) { + model_input_params.selected_token_idxes = + sampling_params.selected_token_idxes; + } + // call model executor forward to get hidden states auto model_output = model_executor_->forward( - input.token_ids, input.positions, kv_caches_, input.input_params); + input.token_ids, input.positions, kv_caches_, model_input_params); if (!model_output.hidden_states.defined()) { return std::nullopt; } @@ -261,7 +267,7 @@ std::optional LLMWorkerImpl::step_internal( torch::Tensor selected_hidden_from_lm_head; if (sampling_params.selected_token_idxes.defined()) { torch::Tensor selected_token_idxes = sampling_params.selected_token_idxes; - if (model_output.hidden_states.defined() && + if (model_output.hidden_states.defined() && selected_token_idxes.defined() && selected_token_idxes.device() != model_output.hidden_states.device()) { selected_token_idxes = selected_token_idxes .to(model_output.hidden_states.device(), diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 161b9200a8..454324ccef 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -142,6 +142,17 @@ struct Options { // enable prefill-only sequence parallel. PROPERTY(bool, enable_prefill_sp) = false; + // Flash Communication 1 (FC1) sequence-parallel optimization. + PROPERTY(bool, enable_flashcomm1) = false; + + PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 1000; + + PROPERTY(int32_t, flashcomm1_min_decode_tokens) = 128; + + PROPERTY(bool, enable_mmrs_fusion) = false; + + PROPERTY(std::string, mmrs_comm_mode) = "aiv"; + // enable returning aux_hidden_states in graph executor output. PROPERTY(bool, enable_graph_aux_hidden_states) = false; diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index a8031edf16..2abef685f9 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -40,6 +40,7 @@ limitations under the License. #include #include "common/device_monitor.h" +#include "core/common/flash_comm1_context.h" #include "common/global_flags.h" #include "common/metrics.h" #include "core/framework/config/beam_search_config.h" @@ -232,6 +233,7 @@ WorkerImpl::WorkerImpl(const ParallelArgs& parallel_args, driver_ = parallel_args.rank() == 0; int32_t tp_size = parallel_args.world_size() / (parallel_args.dp_size() * parallel_args.cp_size()); + dp_driver_ = parallel_args.dp_size() > 1 && parallel_args.rank() % (tp_size * parallel_args.cp_size()) == 0; @@ -1499,6 +1501,15 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, auto tensor_options = torch::dtype(dtype_).device(device_); context_ = ModelContext(parallel_args_, args, quant_args, tensor_options); context_.set_model_id(options_.model_id()); + FlashComm1Options flash_comm1_options; + flash_comm1_options.enable_flashcomm1 = options_.enable_flashcomm1(); + flash_comm1_options.min_prefill_tokens = + options_.flashcomm1_min_prefill_tokens(); + flash_comm1_options.min_decode_tokens = + options_.flashcomm1_min_decode_tokens(); + flash_comm1_options.enable_mmrs_fusion = options_.enable_mmrs_fusion(); + flash_comm1_options.mmrs_comm_mode = options_.mmrs_comm_mode(); + context_.set_flash_comm1_options(flash_comm1_options); // init model, create model executor bool status = this->init_model(context_); diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index be8b0a5e9a..9a7c7ced7d 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -22,11 +22,13 @@ limitations under the License. #include #include +#include "core/common/flash_comm1_context.h" #include "core/framework/kv_cache/kv_cache.h" #include "core/framework/model/model_input_params.h" #include "core/framework/model/model_output.h" #include "core/framework/model_context.h" #include "core/framework/model_loader.h" +#include "core/framework/parallel_state/parallel_args.h" #include "core/layers/common/attention_mask.h" #include "core/layers/common/attention_metadata_builder.h" #include "core/layers/common/lm_head.h" @@ -54,7 +56,9 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { public: explicit Qwen3HybridModelImplBase(const ModelContext& context) : device_(context.get_tensor_options().device()), - model_args_(context.get_model_args()) { + model_args_(context.get_model_args()), + parallel_args_(context.get_parallel_args()), + flash_comm1_options_(context.get_flash_comm1_options()) { auto options = context.get_tensor_options(); auto parallel_args = context.get_parallel_args(); @@ -97,6 +101,13 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { input_params, model_args_.enable_mla(), build_attention_mask(input_params)); + const int32_t num_tokens = static_cast(tokens.size(0)); + const auto& batch_forward_type = input_params.meta.batch_forward_type; + const bool is_prefill_side = batch_forward_type.no_decode(); + FlashComm1Options fc1_options = flash_comm1_options_; + FlashComm1Context fc1_ctx = build_flash_comm1_context( + num_tokens, is_prefill_side, parallel_args_, fc1_options); + torch::Tensor h; if (input_params.embedding.input_embedding.defined()) { h = input_params.embedding.input_embedding; @@ -104,6 +115,10 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { h = embed_tokens_(tokens); } + if (fc1_ctx.is_sequence_sharded()) { + h = shard_sequence(h, fc1_ctx); + } + torch::Tensor mrope_cos_sin; for (const auto& layer : layers_) { mrope_cos_sin = layer->build_mrope_cos_sin(positions); @@ -119,7 +134,8 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { attn_metadata, kv_caches[i], input_params, - mrope_cos_sin); + mrope_cos_sin, + &fc1_ctx); #if defined(USE_NPU) if (input_params.parallel.layer_synchronizer != nullptr && !input_params.parallel.layer_synchronizer->record_event( @@ -130,6 +146,9 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { } auto [hidden_states, residual_out] = norm_->forward(h, residual); h = hidden_states; + if (fc1_ctx.is_sequence_sharded()) { + h = gather_and_unpad_sequence(h, fc1_ctx); + } return ModelOutput(h); } @@ -217,6 +236,8 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { std::vector layers_; int32_t max_seq_len_ = 0; int32_t dp_size_ = 1; + ParallelArgs parallel_args_; + FlashComm1Options flash_comm1_options_; torch::Device device_; torch::ScalarType dtype_ = torch::kFloat; layer::Qwen3NextRMSNorm norm_{nullptr}; diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index 193027ccb8..35df7bd7dc 100644 --- a/xllm/pybind/bind.cpp +++ b/xllm/pybind/bind.cpp @@ -70,6 +70,13 @@ PYBIND11_MODULE(xllm_export, m) { .def_readwrite("enable_chunked_prefill", &Options::enable_chunked_prefill_) .def_readwrite("enable_prefill_sp", &Options::enable_prefill_sp_) + .def_readwrite("enable_flashcomm1", &Options::enable_flashcomm1_) + .def_readwrite("flashcomm1_min_prefill_tokens", + &Options::flashcomm1_min_prefill_tokens_) + .def_readwrite("flashcomm1_min_decode_tokens", + &Options::flashcomm1_min_decode_tokens_) + .def_readwrite("enable_mmrs_fusion", &Options::enable_mmrs_fusion_) + .def_readwrite("mmrs_comm_mode", &Options::mmrs_comm_mode_) .def_readwrite("master_node_addr", &Options::master_node_addr_) .def_readwrite("nnodes", &Options::nnodes_) .def_readwrite("node_rank", &Options::node_rank_) diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index a923f5f278..13a5b076ca 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -26,6 +26,7 @@ limitations under the License. #include #include "api_service/api_service.h" +#include "core/common/global_flags.h" #include "core/common/instance_name.h" #include "core/common/metrics.h" #include "core/common/options.h" @@ -157,6 +158,11 @@ Options create_options(const std::string& instance_name, bool is_local) { .expert_parallel_degree(eplb_config.expert_parallel_degree()) .enable_chunked_prefill(scheduler_config.enable_chunked_prefill()) .enable_prefill_sp(parallel_config.enable_prefill_sp()) + .enable_flashcomm1(FLAGS_enable_flashcomm1) + .flashcomm1_min_prefill_tokens(FLAGS_flashcomm1_min_prefill_tokens) + .flashcomm1_min_decode_tokens(FLAGS_flashcomm1_min_decode_tokens) + .enable_mmrs_fusion(FLAGS_enable_mmrs_fusion) + .mmrs_comm_mode(FLAGS_mmrs_comm_mode) .master_node_addr(distributed_config.master_node_addr()) .instance_role(InstanceRole(disagg_pd_config.instance_role())) .transfer_listen_port( From 17a96485c75186315ed7779fb6fbb475e776a7bc Mon Sep 17 00:00:00 2001 From: Cooofish <2716856597@qq.com> Date: Tue, 7 Jul 2026 05:07:32 +0000 Subject: [PATCH 2/6] bugfix: make flashcomm1 pr build reliably. - Infer NPU Python, PyTorch, and torch_npu paths during CMake configure when the environment is unset. - Propagate CARGO_HOME into Rust custom commands and align the workspace Cargo mirror name with the root Cargo config. - Remove stale selected-gather and Device::type_str() remnants after the FC1 context refactor. - Keep FC1 default-off and raise the activation threshold to long-prefill TP=8 cases. --- CMakeLists.txt | 46 ++++++++++++++++++- cmake/FindRust.cmake | 10 ++-- cmake/cargo_library.cmake | 4 +- cmake/cargo_shared_library.cmake | 4 +- third_party/xllm_atb_layers | 2 +- third_party/xllm_ops | 2 +- xllm/core/common/flash_comm1_context.cpp | 36 ++++++++------- xllm/core/common/flash_comm1_context.h | 17 ++++++- xllm/core/common/global_flags.h | 2 - xllm/core/common/options.cpp | 1 - xllm/core/common/options.h | 4 +- xllm/core/distributed_runtime/master.cpp | 3 -- xllm/core/layers/common/dense_mlp.cpp | 13 +++--- xllm/core/layers/common/dense_mlp.h | 3 +- xllm/core/layers/common/linear.cpp | 4 +- xllm/core/layers/common/linear.h | 3 +- .../npu_torch/qwen3_gated_delta_net_base.cpp | 12 +---- .../npu_torch/qwen3_gated_delta_net_base.h | 6 --- .../layers/npu_torch/qwen3_next_attention.cpp | 15 ++---- .../layers/npu_torch/qwen3_next_attention.h | 9 +--- .../qwen3_next_hybrid_decoder_layer_base.cpp | 26 ++--------- .../qwen3_next_hybrid_decoder_layer_base.h | 21 +-------- xllm/core/runtime/llm_worker_impl.cpp | 4 -- xllm/core/runtime/options.h | 4 +- xllm/core/runtime/worker_impl.cpp | 2 - xllm/models/llm/qwen3_next_hybrid_base.h | 4 +- xllm/pybind/bind.cpp | 2 - xllm/xllm.cpp | 1 - 28 files changed, 119 insertions(+), 141 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9090501db7..f4cfa3cd0b 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -334,9 +334,53 @@ find_package(benchmark CONFIG REQUIRED) find_package(nlohmann_json CONFIG REQUIRED) find_package(OpenCV CONFIG REQUIRED) find_package(FFMPEG REQUIRED) -find_package(Python COMPONENTS Development REQUIRED) +find_package(Python COMPONENTS Interpreter Development REQUIRED) find_package(pybind11 CONFIG REQUIRED) +if(USE_NPU) + function(xllm_set_env_from_python VAR PYTHON_CODE) + if("$ENV{${VAR}}" STREQUAL "") + execute_process( + COMMAND ${Python_EXECUTABLE} -c "${PYTHON_CODE}" + OUTPUT_VARIABLE XLLM_INFERRED_ENV_VALUE + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_VARIABLE XLLM_INFERRED_ENV_ERROR + RESULT_VARIABLE XLLM_INFERRED_ENV_RESULT + ) + if(NOT XLLM_INFERRED_ENV_RESULT EQUAL 0 OR "${XLLM_INFERRED_ENV_VALUE}" STREQUAL "") + message(FATAL_ERROR "Failed to infer ${VAR}: ${XLLM_INFERRED_ENV_ERROR}") + endif() + set(ENV{${VAR}} "${XLLM_INFERRED_ENV_VALUE}") + message(STATUS "Inferred ${VAR}: $ENV{${VAR}}") + endif() + endfunction() + + if("$ENV{PYTHON_INCLUDE_PATH}" STREQUAL "") + set(ENV{PYTHON_INCLUDE_PATH} "${Python_INCLUDE_DIRS}") + endif() + if("$ENV{PYTHON_LIB_PATH}" STREQUAL "") + execute_process( + COMMAND ${Python_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR') or '')" + OUTPUT_VARIABLE XLLM_PYTHON_LIB_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + set(ENV{PYTHON_LIB_PATH} "${XLLM_PYTHON_LIB_PATH}") + endif() + xllm_set_env_from_python(PYTORCH_INSTALL_PATH "import importlib.util; spec = importlib.util.find_spec('torch'); print(spec.submodule_search_locations[0])") + xllm_set_env_from_python(PYTORCH_NPU_INSTALL_PATH "import importlib.util; spec = importlib.util.find_spec('torch_npu'); print(spec.submodule_search_locations[0])") + if(NOT EXISTS "$ENV{PYTORCH_NPU_INSTALL_PATH}/include/torch_npu/torch_npu.h" + AND EXISTS "/usr/local/libtorch_npu/include/torch_npu/torch_npu.h") + set(ENV{PYTORCH_NPU_INSTALL_PATH} "/usr/local/libtorch_npu") + message(STATUS "Using PYTORCH_NPU_INSTALL_PATH with compatibility header: $ENV{PYTORCH_NPU_INSTALL_PATH}") + endif() + if("$ENV{LIBTORCH_ROOT}" STREQUAL "") + set(ENV{LIBTORCH_ROOT} "$ENV{PYTORCH_INSTALL_PATH}") + endif() + if("$ENV{NPU_HOME_PATH}" STREQUAL "" AND NOT "$ENV{ASCEND_HOME_PATH}" STREQUAL "") + set(ENV{NPU_HOME_PATH} "$ENV{ASCEND_HOME_PATH}") + endif() +endif() + if (USE_CXX11_ABI) # only use jemalloc if using the new C++-11 ABI find_package(Jemalloc) diff --git a/cmake/FindRust.cmake b/cmake/FindRust.cmake index 8c771250dd..b3e0f995fb 100644 --- a/cmake/FindRust.cmake +++ b/cmake/FindRust.cmake @@ -10,12 +10,10 @@ else() set(USER_HOME "$ENV{HOME}") endif() -if(NOT DEFINED CARGO_HOME) - if("$ENV{CARGO_HOME}" STREQUAL "") - set(CARGO_HOME "${USER_HOME}/.cargo") - else() - set(CARGO_HOME "$ENV{CARGO_HOME}") - endif() +if(NOT "$ENV{CARGO_HOME}" STREQUAL "") + set(CARGO_HOME "$ENV{CARGO_HOME}" CACHE PATH "Rust Cargo Home" FORCE) +elseif(NOT DEFINED CARGO_HOME) + set(CARGO_HOME "${USER_HOME}/.cargo") endif() # Find cargo executable diff --git a/cmake/cargo_library.cmake b/cmake/cargo_library.cmake index bf34819525..fadc5725ce 100644 --- a/cmake/cargo_library.cmake +++ b/cmake/cargo_library.cmake @@ -65,7 +65,9 @@ function(cargo_library) file(GLOB_RECURSE LIB_SOURCES "*.rs") - set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") + set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env + "CARGO_HOME=${CARGO_HOME}" + "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") # build the library target with cargo set(STATIC_LIB_NAME diff --git a/cmake/cargo_shared_library.cmake b/cmake/cargo_shared_library.cmake index 8bac49f016..d92e7f7ab1 100644 --- a/cmake/cargo_shared_library.cmake +++ b/cmake/cargo_shared_library.cmake @@ -65,7 +65,9 @@ function(cargo_shared_library) file(GLOB_RECURSE LIB_SOURCES "*.rs") - set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") + set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env + "CARGO_HOME=${CARGO_HOME}" + "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") # build the library target with cargo set(SHARED_LIB_NAME diff --git a/third_party/xllm_atb_layers b/third_party/xllm_atb_layers index 4b926ee633..5b50a7e167 160000 --- a/third_party/xllm_atb_layers +++ b/third_party/xllm_atb_layers @@ -1 +1 @@ -Subproject commit 4b926ee633ed11677890d7e2ddbad7eb6dd249f2 +Subproject commit 5b50a7e167267d4482c665717f25e1ef066a8e86 diff --git a/third_party/xllm_ops b/third_party/xllm_ops index bbbdfe3818..e7d82584a7 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit bbbdfe3818149b6d4d725e1fbf44149c4c01a1a6 +Subproject commit e7d82584a7c3043b4b5a08f60643e04537771147 diff --git a/xllm/core/common/flash_comm1_context.cpp b/xllm/core/common/flash_comm1_context.cpp index 3ea1ec1791..7b47220867 100644 --- a/xllm/core/common/flash_comm1_context.cpp +++ b/xllm/core/common/flash_comm1_context.cpp @@ -21,9 +21,6 @@ limitations under the License. #include "common/global_flags.h" #include "framework/parallel_state/parallel_state.h" -#if defined(USE_NPU) -#include "platform/device.h" -#endif DEFINE_bool(enable_flashcomm1, false, @@ -31,13 +28,9 @@ DEFINE_bool(enable_flashcomm1, "for tensor parallel inference on NPU."); DEFINE_int32(flashcomm1_min_prefill_tokens, - 1000, + 8192, "Minimum prefill token count to activate FC1."); -DEFINE_int32(flashcomm1_min_decode_tokens, - 128, - "Minimum decode batch token count to activate FC1."); - DEFINE_bool(enable_mmrs_fusion, false, "Enable Matmul+ReduceScatter fusion kernel for FC1."); @@ -52,6 +45,8 @@ namespace xllm { namespace { constexpr int32_t kFc1LocalTokenAlignment = 16; +constexpr int32_t kFc1MinTpSize = 8; +thread_local const FlashComm1Context* current_flash_comm1_context = nullptr; int32_t round_up_to_multiple(int32_t value, int32_t multiple) { CHECK_GT(multiple, 0); @@ -92,6 +87,20 @@ torch::Tensor pad_rows_by_copy(const torch::Tensor& input, } // namespace +FlashComm1ContextScope::FlashComm1ContextScope( + const FlashComm1Context* ctx) + : previous_(current_flash_comm1_context) { + current_flash_comm1_context = ctx; +} + +FlashComm1ContextScope::~FlashComm1ContextScope() { + current_flash_comm1_context = previous_; +} + +const FlashComm1Context* get_current_flash_comm1_context() { + return current_flash_comm1_context; +} + int32_t FlashComm1Context::local_num_tokens_for_rank(int32_t rank) const { CHECK_GE(rank, 0); CHECK_LT(rank, tp_world_size); @@ -133,11 +142,7 @@ FlashComm1Context build_flash_comm1_context( return ctx; } -#if defined(USE_NPU) - if (Device::type_str() != "npu") { - return ctx; - } -#else +#if !defined(USE_NPU) return ctx; #endif @@ -145,7 +150,7 @@ FlashComm1Context build_flash_comm1_context( return ctx; } - if (actual_tp_size <= 1) { + if (actual_tp_size < kFc1MinTpSize) { return ctx; } @@ -156,7 +161,7 @@ FlashComm1Context build_flash_comm1_context( int32_t threshold = options.min_prefill_tokens; - if (num_tokens <= threshold) { + if (num_tokens < threshold) { return ctx; } @@ -185,7 +190,6 @@ FlashComm1Context build_flash_comm1_context(int32_t num_tokens, FlashComm1Options options; options.enable_flashcomm1 = FLAGS_enable_flashcomm1; options.min_prefill_tokens = FLAGS_flashcomm1_min_prefill_tokens; - options.min_decode_tokens = FLAGS_flashcomm1_min_decode_tokens; options.enable_mmrs_fusion = FLAGS_enable_mmrs_fusion; options.mmrs_comm_mode = FLAGS_mmrs_comm_mode; return build_flash_comm1_context( diff --git a/xllm/core/common/flash_comm1_context.h b/xllm/core/common/flash_comm1_context.h index 73fba5aee8..bb0ccf4435 100644 --- a/xllm/core/common/flash_comm1_context.h +++ b/xllm/core/common/flash_comm1_context.h @@ -61,12 +61,25 @@ struct FlashComm1Context { struct FlashComm1Options { bool enable_flashcomm1 = false; - int32_t min_prefill_tokens = 1000; - int32_t min_decode_tokens = 128; + int32_t min_prefill_tokens = 8192; bool enable_mmrs_fusion = false; std::string mmrs_comm_mode = "aiv"; }; +class FlashComm1ContextScope { + public: + explicit FlashComm1ContextScope(const FlashComm1Context* ctx); + ~FlashComm1ContextScope(); + + FlashComm1ContextScope(const FlashComm1ContextScope&) = delete; + FlashComm1ContextScope& operator=(const FlashComm1ContextScope&) = delete; + + private: + const FlashComm1Context* previous_; +}; + +const FlashComm1Context* get_current_flash_comm1_context(); + FlashComm1Context build_flash_comm1_context( int32_t num_tokens, bool is_prefill, diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index a5f17ef488..85adfc9c9b 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -393,8 +393,6 @@ DECLARE_bool(enable_flashcomm1); DECLARE_int32(flashcomm1_min_prefill_tokens); -DECLARE_int32(flashcomm1_min_decode_tokens); - DECLARE_bool(enable_mmrs_fusion); DECLARE_string(mmrs_comm_mode); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index c0edda7903..62d74b6a09 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -59,7 +59,6 @@ std::string Options::to_string() const { << ", enable_flashcomm1: " << enable_flashcomm1() << ", flashcomm1_min_prefill_tokens: " << flashcomm1_min_prefill_tokens() - << ", flashcomm1_min_decode_tokens: " << flashcomm1_min_decode_tokens() << ", enable_mmrs_fusion: " << enable_mmrs_fusion() << ", mmrs_comm_mode: " << mmrs_comm_mode() << ", master_node_addr: " << master_node_addr().value_or("null") diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index b5eb791193..c0cf11f40e 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -125,9 +125,7 @@ class Options { // Flash Communication 1 (FC1) sequence-parallel optimization. PROPERTY(bool, enable_flashcomm1) = false; - PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 1000; - - PROPERTY(int32_t, flashcomm1_min_decode_tokens) = 128; + PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 8192; PROPERTY(bool, enable_mmrs_fusion) = false; diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index 89737184dc..dbbe63fe00 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -243,7 +243,6 @@ Master::Master(const Options& options, EngineType type) .enable_flashcomm1(options_.enable_flashcomm1()) .flashcomm1_min_prefill_tokens( options_.flashcomm1_min_prefill_tokens()) - .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) .enable_mmrs_fusion(options_.enable_mmrs_fusion()) .mmrs_comm_mode(options_.mmrs_comm_mode()) .npu_kernel_backend(options_.npu_kernel_backend()) @@ -322,7 +321,6 @@ Master::Master(const Options& options, EngineType type) .enable_flashcomm1(options_.enable_flashcomm1()) .flashcomm1_min_prefill_tokens( options_.flashcomm1_min_prefill_tokens()) - .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) .enable_mmrs_fusion(options_.enable_mmrs_fusion()) .mmrs_comm_mode(options_.mmrs_comm_mode()) .cp_size(options.cp_size()) @@ -383,7 +381,6 @@ Master::Master(const Options& options, EngineType type) .enable_flashcomm1(options_.enable_flashcomm1()) .flashcomm1_min_prefill_tokens( options_.flashcomm1_min_prefill_tokens()) - .flashcomm1_min_decode_tokens(options_.flashcomm1_min_decode_tokens()) .enable_mmrs_fusion(options_.enable_mmrs_fusion()) .mmrs_comm_mode(options_.mmrs_comm_mode()) .cp_size(options_.cp_size()) diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index cd2232aa9d..760ba4db0b 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -19,7 +19,6 @@ limitations under the License. #include "common/flash_comm1_context.h" #include "kernels/ops_api.h" -#include "platform/device.h" namespace xllm { namespace layer { @@ -94,8 +93,8 @@ DenseMLPImpl::DenseMLPImpl(int64_t hidden_size, down_proj_extra_args)); } -torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states, - const FlashComm1Context* fc1_ctx) { +torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states) { + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { @@ -110,18 +109,20 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states, << "FC1 MMRS callsite DenseMLP.down_proj(smoothquant): input=" << gate_up.sizes(); return down_proj_->forward( - gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } return down_proj_->forward(gate_up); } torch::Tensor output; - if (Device::type_str() != "npu") { +#if !defined(USE_NPU) + { int64_t batch_size = gate_up.sizes()[0]; output = torch::empty( {batch_size, intermediate_size_ / process_group_->world_size()}, gate_up.options()); } +#endif act_->forward(gate_up, output); @@ -129,7 +130,7 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states, LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite DenseMLP.down_proj: input=" << output.sizes(); return down_proj_->forward( - output, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + output, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } return down_proj_->forward(output); } diff --git a/xllm/core/layers/common/dense_mlp.h b/xllm/core/layers/common/dense_mlp.h index d2cdeff310..291f0d257a 100644 --- a/xllm/core/layers/common/dense_mlp.h +++ b/xllm/core/layers/common/dense_mlp.h @@ -45,8 +45,7 @@ class DenseMLPImpl : public torch::nn::Module { const std::string& module_prefix = "", double swiglu_limit = 0.0); - torch::Tensor forward(torch::Tensor hidden_states, - const FlashComm1Context* fc1_ctx = nullptr); + torch::Tensor forward(torch::Tensor hidden_states); void load_state_dict(const StateDict& state_dict); void load_state_dict(const StateDict& state_dict, diff --git a/xllm/core/layers/common/linear.cpp b/xllm/core/layers/common/linear.cpp index 50bcc1ebc3..7f2420889b 100644 --- a/xllm/core/layers/common/linear.cpp +++ b/xllm/core/layers/common/linear.cpp @@ -1571,8 +1571,8 @@ torch::Tensor RowParallelLinearImpl::mmrs_weight_transposed() const { } torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, - RowParallelReduceMode reduce_mode, - const FlashComm1Context* fc1_ctx) { + RowParallelReduceMode reduce_mode) { + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); auto bias = bias_.defined() && rank_ == 0 ? std::optional(bias_) : std::nullopt; diff --git a/xllm/core/layers/common/linear.h b/xllm/core/layers/common/linear.h index 35db282059..00c1361a64 100644 --- a/xllm/core/layers/common/linear.h +++ b/xllm/core/layers/common/linear.h @@ -261,8 +261,7 @@ class RowParallelLinearImpl : public torch::nn::Module { torch::Tensor forward(torch::Tensor input); torch::Tensor forward(torch::Tensor input, - RowParallelReduceMode reduce_mode, - const FlashComm1Context* fc1_ctx); + RowParallelReduceMode reduce_mode); // load the weight from the checkpoint void load_state_dict(const StateDict& state_dict); diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index 06aea7069b..e969b5a62e 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -533,15 +533,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params) { - return forward(hidden_states, attn_metadata, kv_cache, input_params, nullptr); -} - -torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const FlashComm1Context* fc1_ctx) { + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { h = gather_sequence(hidden_states, *fc1_ctx); @@ -927,7 +919,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( << "FC1 MMRS callsite Qwen3GatedDeltaNet.o_proj: input=" << rearranged_norm.sizes(); return o_proj_->forward( - rearranged_norm, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + rearranged_norm, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } return o_proj_->forward(rearranged_norm); } diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h index a8350356d7..9dbf0bcfaf 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h @@ -51,12 +51,6 @@ class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { KVCache& kv_cache, const ModelInputParams& input_params); - torch::Tensor forward(const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const FlashComm1Context* fc1_ctx); - protected: virtual std::pair project_decode_inputs( const torch::Tensor& hidden_states) = 0; diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index 6134cb1d86..646c4d564d 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -147,16 +147,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( const AttentionMetadata& attn_metadata, KVCache& kv_cache, const torch::Tensor& mrope_cos_sin) { - return forward(positions, hidden_states, attn_metadata, kv_cache, mrope_cos_sin, nullptr); -} - -torch::Tensor Qwen3NextAttentionImpl::forward( - const torch::Tensor& positions, - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const torch::Tensor& mrope_cos_sin, - const FlashComm1Context* fc1_ctx) { + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { @@ -193,7 +184,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( << "FC1 MMRS callsite Qwen3NextAttention.o_proj(mrope): input=" << out.sizes(); return o_proj_->forward( - out, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + out, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } return o_proj_->forward(out); } @@ -230,7 +221,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( << "FC1 MMRS callsite Qwen3NextAttention.o_proj: input=" << out.sizes(); return o_proj_->forward( - out, row_parallel_reduce_mode_for_fc1(*fc1_ctx), fc1_ctx); + out, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } return o_proj_->forward(out); } diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.h b/xllm/core/layers/npu_torch/qwen3_next_attention.h index dbfd34e37c..f6c3ff1a5a 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.h +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.h @@ -43,18 +43,11 @@ class Qwen3NextAttentionImpl : public torch::nn::Module { const torch::TensorOptions& options, int32_t layer_id); -torch::Tensor forward(const torch::Tensor& positions, - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const torch::Tensor& mrope_cos_sin); - torch::Tensor forward(const torch::Tensor& positions, const torch::Tensor& hidden_states, const AttentionMetadata& attn_metadata, KVCache& kv_cache, - const torch::Tensor& mrope_cos_sin, - const FlashComm1Context* fc1_ctx); + const torch::Tensor& mrope_cos_sin); torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const; diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp index b3a0ea4ecd..81426acee9 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp @@ -116,25 +116,7 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( KVCache& kv_cache, const ModelInputParams& input_params, const torch::Tensor& mrope_cos_sin) { - return forward(x, - residual, - positions, - attn_metadata, - kv_cache, - input_params, - mrope_cos_sin, - nullptr); -} - -torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( - torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin, - const FlashComm1Context* fc1_ctx) { + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); if (!residual.has_value()) { residual = x; x = std::get<0>(input_norm_->forward(x)); @@ -153,10 +135,10 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( if (attention_) { x = attention_->forward( - positions, x, attn_metadata, kv_cache, mrope_cos_sin, fc1_ctx); + positions, x, attn_metadata, kv_cache, mrope_cos_sin); } else { x = linear_attention_->forward( - x, attn_metadata, kv_cache, input_params, fc1_ctx); + x, attn_metadata, kv_cache, input_params); } // Before post_norm, ensure residual shape matches x shape @@ -173,7 +155,7 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( if (moe_mlp_) { x = moe_mlp_(x, input_params); } else { - x = mlp_->forward(x, fc1_ctx); + x = mlp_->forward(x); } return x; diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h index f50b4f1b2f..e431a1698a 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h @@ -38,21 +38,13 @@ class Qwen3HybridDecoderLayerModule : public torch::nn::Module { public: virtual void load_state_dict(const StateDict& state_dict) = 0; virtual void verify_loaded_weights(const std::string& prefix) const = 0; -virtual torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) = 0; virtual torch::Tensor forward(torch::Tensor& x, std::optional& residual, torch::Tensor& positions, const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin, - const FlashComm1Context* fc1_ctx) = 0; + const torch::Tensor& mrope_cos_sin = {}) = 0; virtual torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const { return {}; @@ -73,22 +65,13 @@ class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule { void verify_loaded_weights(const std::string& prefix) const override; -torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) override; - torch::Tensor forward(torch::Tensor& x, std::optional& residual, torch::Tensor& positions, const AttentionMetadata& attn_metadata, KVCache& kv_cache, const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin, - const FlashComm1Context* fc1_ctx) override; + const torch::Tensor& mrope_cos_sin = {}) override; torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const override; diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index c567317fe7..439da3aa55 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -251,10 +251,6 @@ std::optional LLMWorkerImpl::step_internal( } ModelInputParams model_input_params = input.input_params; - if (sampling_params.selected_token_idxes.defined()) { - model_input_params.selected_token_idxes = - sampling_params.selected_token_idxes; - } // call model executor forward to get hidden states auto model_output = model_executor_->forward( diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 454324ccef..c732ac302a 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -145,9 +145,7 @@ struct Options { // Flash Communication 1 (FC1) sequence-parallel optimization. PROPERTY(bool, enable_flashcomm1) = false; - PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 1000; - - PROPERTY(int32_t, flashcomm1_min_decode_tokens) = 128; + PROPERTY(int32_t, flashcomm1_min_prefill_tokens) = 8192; PROPERTY(bool, enable_mmrs_fusion) = false; diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 2abef685f9..87bfe3ca6e 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -1505,8 +1505,6 @@ bool WorkerImpl::init_model(const std::string& model_weights_path, flash_comm1_options.enable_flashcomm1 = options_.enable_flashcomm1(); flash_comm1_options.min_prefill_tokens = options_.flashcomm1_min_prefill_tokens(); - flash_comm1_options.min_decode_tokens = - options_.flashcomm1_min_decode_tokens(); flash_comm1_options.enable_mmrs_fusion = options_.enable_mmrs_fusion(); flash_comm1_options.mmrs_comm_mode = options_.mmrs_comm_mode(); context_.set_flash_comm1_options(flash_comm1_options); diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index 9a7c7ced7d..0ba313813e 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -107,6 +107,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { FlashComm1Options fc1_options = flash_comm1_options_; FlashComm1Context fc1_ctx = build_flash_comm1_context( num_tokens, is_prefill_side, parallel_args_, fc1_options); + FlashComm1ContextScope fc1_scope(&fc1_ctx); torch::Tensor h; if (input_params.embedding.input_embedding.defined()) { @@ -134,8 +135,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { attn_metadata, kv_caches[i], input_params, - mrope_cos_sin, - &fc1_ctx); + mrope_cos_sin); #if defined(USE_NPU) if (input_params.parallel.layer_synchronizer != nullptr && !input_params.parallel.layer_synchronizer->record_event( diff --git a/xllm/pybind/bind.cpp b/xllm/pybind/bind.cpp index 35df7bd7dc..31efe619b6 100644 --- a/xllm/pybind/bind.cpp +++ b/xllm/pybind/bind.cpp @@ -73,8 +73,6 @@ PYBIND11_MODULE(xllm_export, m) { .def_readwrite("enable_flashcomm1", &Options::enable_flashcomm1_) .def_readwrite("flashcomm1_min_prefill_tokens", &Options::flashcomm1_min_prefill_tokens_) - .def_readwrite("flashcomm1_min_decode_tokens", - &Options::flashcomm1_min_decode_tokens_) .def_readwrite("enable_mmrs_fusion", &Options::enable_mmrs_fusion_) .def_readwrite("mmrs_comm_mode", &Options::mmrs_comm_mode_) .def_readwrite("master_node_addr", &Options::master_node_addr_) diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 13a5b076ca..31c96cef1a 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -160,7 +160,6 @@ Options create_options(const std::string& instance_name, bool is_local) { .enable_prefill_sp(parallel_config.enable_prefill_sp()) .enable_flashcomm1(FLAGS_enable_flashcomm1) .flashcomm1_min_prefill_tokens(FLAGS_flashcomm1_min_prefill_tokens) - .flashcomm1_min_decode_tokens(FLAGS_flashcomm1_min_decode_tokens) .enable_mmrs_fusion(FLAGS_enable_mmrs_fusion) .mmrs_comm_mode(FLAGS_mmrs_comm_mode) .master_node_addr(distributed_config.master_node_addr()) From 48672c1096425de9bc89c4e265d09472deec5400 Mon Sep 17 00:00:00 2001 From: Han Teng <165119118+Cooofish@users.noreply.github.com> Date: Tue, 7 Jul 2026 14:18:47 +0800 Subject: [PATCH 3/6] refactor: simplify flashcomm1 mmrs wrapper. - Keep the wrapper focused on validating inputs, resolving the HCCL group, selecting comm_mode, and invoking torch_npu npu_mm_reduce_scatter_base. - Return an empty tensor on unsupported inputs so the existing row-parallel matmul plus reduce_scatter fallback remains owned by the caller. - Remove the duplicate local torch::matmul fallback and reject-reason string builder to make the FC1 MMRS path easier to review and maintain. --- .../kernels/npu/matmul_reduce_scatter.cpp | 85 ++++++++----------- 1 file changed, 36 insertions(+), 49 deletions(-) diff --git a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp index af686290c9..938ac1bf73 100644 --- a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp +++ b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp @@ -18,73 +18,69 @@ limitations under the License. #include #include -#include - #include "core/framework/parallel_state/process_group.h" namespace xllm::kernel::npu { namespace { -std::string matmul_reduce_scatter_reject_reason( +bool can_call_torch_npu_mmrs( const torch::Tensor& a, const torch::Tensor& b, const std::optional& bias, const std::optional& output, ProcessGroup* process_group) { if (process_group == nullptr) { - return "process_group is null"; + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: process_group is null."; + return false; } if (!output.has_value() || !output->defined()) { - return "output tensor is missing"; + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: output tensor is missing."; + return false; } if (a.dim() != 2 || b.dim() != 2 || output->dim() != 2) { - std::ostringstream oss; - oss << "expected 2D tensors, got a_dim=" << a.dim() - << ", b_dim=" << b.dim() << ", output_dim=" << output->dim(); - return oss.str(); + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: expected 2D tensors, got a_dim=" + << a.dim() << ", b_dim=" << b.dim() + << ", output_dim=" << output->dim(); + return false; } if (a.scalar_type() != at::kHalf && a.scalar_type() != at::kBFloat16) { - std::ostringstream oss; - oss << "unsupported input dtype=" << a.scalar_type(); - return oss.str(); + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: unsupported input dtype=" + << a.scalar_type(); + return false; } if (a.scalar_type() != b.scalar_type() || a.scalar_type() != output->scalar_type()) { - std::ostringstream oss; - oss << "dtype mismatch: a=" << a.scalar_type() - << ", b=" << b.scalar_type() + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: dtype mismatch. a=" + << a.scalar_type() << ", b=" << b.scalar_type() << ", output=" << output->scalar_type(); - return oss.str(); + return false; } if (bias.has_value() && bias->defined() && bias->scalar_type() != a.scalar_type()) { - std::ostringstream oss; - oss << "bias dtype mismatch: bias=" << bias->scalar_type() - << ", input=" << a.scalar_type(); - return oss.str(); + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: bias dtype mismatch. bias=" + << bias->scalar_type() << ", input=" << a.scalar_type(); + return false; } if (a.size(1) != b.size(0)) { - std::ostringstream oss; - oss << "matmul K mismatch: a=" << a.sizes() << ", b=" << b.sizes(); - return oss.str(); + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: matmul K mismatch. a=" << a.sizes() + << ", b=" << b.sizes(); + return false; } if (output->size(1) != b.size(1)) { - std::ostringstream oss; - oss << "output N mismatch: output=" << output->sizes() + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: output N mismatch. output=" + << output->sizes() << ", b=" << b.sizes(); - return oss.str(); - } - return ""; -} - -torch::Tensor matmul_kn(const torch::Tensor& a, - const torch::Tensor& b, - const std::optional& bias) { - torch::Tensor out = torch::matmul(a, b); - if (bias.has_value() && bias->defined()) { - out = out + bias.value(); + return false; } - return out; + return true; } bool should_use_ai_cpu_for_aiv_risky_shape(const torch::Tensor& a, @@ -121,17 +117,8 @@ torch::Tensor matmul_reduce_scatter( int64_t comm_turn, int64_t stream_mode, const std::string& comm_mode) { - const std::string reject_reason = - matmul_reduce_scatter_reject_reason(a, b, bias, output, process_group); - if (!reject_reason.empty()) { - LOG_FIRST_N(WARNING, 8) - << "FC1 MMRS torch_npu skipped: " << reject_reason - << "; fallback to matmul + reduce_scatter path. a=" << a.sizes() - << ", b=" << b.sizes() - << ", output=" - << (output.has_value() && output->defined() ? output->sizes() - : c10::IntArrayRef{}); - return matmul_kn(a, b, bias); + if (!can_call_torch_npu_mmrs(a, b, bias, output, process_group)) { + return torch::Tensor(); } std::string group = process_group->hccl_comm_name(/*init_comm=*/true); @@ -139,7 +126,7 @@ torch::Tensor matmul_reduce_scatter( LOG_FIRST_N(WARNING, 8) << "FC1 MMRS torch_npu skipped: HCCL group name is empty; fallback to " "matmul + reduce_scatter path."; - return matmul_kn(a, b, bias); + return torch::Tensor(); } std::string effective_comm_mode = comm_mode; From 5f0a4e0f8668193c839090311643e191d89f990e Mon Sep 17 00:00:00 2001 From: Cooofish <2716856597@qq.com> Date: Tue, 7 Jul 2026 15:38:27 +0800 Subject: [PATCH 4/6] bugfix: address flashcomm1 review comments. - Keep FlashComm1Context as a plain data struct and move sequence-shard helpers to free functions. - Reuse the shared FC1 row-padding helper from linear callsites instead of duplicating it. - Add defensive defined checks before inspecting MMRS input tensor metadata. - Update FC1 callsites to use the free helper while preserving the existing config and MoE reduction fixes. --- xllm/core/common/flash_comm1_context.cpp | 95 ++++++++++--------- xllm/core/common/flash_comm1_context.h | 83 ++++++++-------- .../kernels/npu/matmul_reduce_scatter.cpp | 6 ++ xllm/core/layers/common/dense_mlp.cpp | 6 +- xllm/core/layers/common/linear.cpp | 25 ++--- .../npu_torch/qwen3_gated_delta_net_base.cpp | 4 +- .../layers/npu_torch/qwen3_next_attention.cpp | 6 +- .../qwen3_next_hybrid_decoder_layer_base.cpp | 6 +- xllm/models/llm/qwen3_next_hybrid_base.h | 4 +- 9 files changed, 120 insertions(+), 115 deletions(-) diff --git a/xllm/core/common/flash_comm1_context.cpp b/xllm/core/common/flash_comm1_context.cpp index 7b47220867..a11fd01c20 100644 --- a/xllm/core/common/flash_comm1_context.cpp +++ b/xllm/core/common/flash_comm1_context.cpp @@ -54,37 +54,22 @@ int32_t round_up_to_multiple(int32_t value, int32_t multiple) { return remainder == 0 ? value : value + multiple - remainder; } -int32_t local_num_tokens_for_rank(int32_t num_tokens, - int32_t world_size, - int32_t rank) { +int32_t local_num_tokens_for_rank_impl(int32_t num_tokens, + int32_t world_size, + int32_t rank) { const int32_t base = num_tokens / world_size; const int32_t remainder = num_tokens % world_size; return base + (rank < remainder ? 1 : 0); } -int64_t shard_start_for_rank(int32_t num_tokens, - int32_t world_size, - int32_t rank) { +int64_t shard_start_for_rank_impl(int32_t num_tokens, + int32_t world_size, + int32_t rank) { const int32_t base = num_tokens / world_size; const int32_t remainder = num_tokens % world_size; return static_cast(rank) * base + std::min(rank, remainder); } -torch::Tensor pad_rows_by_copy(const torch::Tensor& input, - int64_t padded_rows) { - CHECK_GE(padded_rows, input.size(0)); - if (padded_rows == input.size(0)) { - return input; - } - - auto output_shape = input.sizes().vec(); - output_shape[0] = padded_rows; - auto output = torch::empty(output_shape, input.options()); - output.slice(0, 0, input.size(0)).copy_(input); - output.slice(0, input.size(0), padded_rows).zero_(); - return output; -} - } // namespace FlashComm1ContextScope::FlashComm1ContextScope( @@ -101,27 +86,47 @@ const FlashComm1Context* get_current_flash_comm1_context() { return current_flash_comm1_context; } -int32_t FlashComm1Context::local_num_tokens_for_rank(int32_t rank) const { +bool is_sequence_sharded(const FlashComm1Context& ctx) { + return ctx.enabled && ctx.tp_world_size > 1; +} + +int32_t local_num_tokens_for_rank(const FlashComm1Context& ctx, int32_t rank) { CHECK_GE(rank, 0); - CHECK_LT(rank, tp_world_size); - return xllm::local_num_tokens_for_rank( - original_num_tokens, tp_world_size, rank); + CHECK_LT(rank, ctx.tp_world_size); + return local_num_tokens_for_rank_impl( + ctx.original_num_tokens, ctx.tp_world_size, rank); } -std::vector FlashComm1Context::token_num_list() const { - std::vector token_nums(tp_world_size); - for (int32_t rank = 0; rank < tp_world_size; ++rank) { - token_nums[rank] = local_num_tokens_for_rank(rank); +std::vector token_num_list(const FlashComm1Context& ctx) { + std::vector token_nums(ctx.tp_world_size); + for (int32_t rank = 0; rank < ctx.tp_world_size; ++rank) { + token_nums[rank] = local_num_tokens_for_rank(ctx, rank); } return token_nums; } -int64_t FlashComm1Context::get_shard_start() const { - return shard_start_for_rank(original_num_tokens, tp_world_size, tp_rank); +int64_t get_shard_start(const FlashComm1Context& ctx) { + return shard_start_for_rank_impl( + ctx.original_num_tokens, ctx.tp_world_size, ctx.tp_rank); +} + +int64_t get_shard_end(const FlashComm1Context& ctx) { + return get_shard_start(ctx) + ctx.local_num_tokens; } -int64_t FlashComm1Context::get_shard_end() const { - return get_shard_start() + local_num_tokens; +torch::Tensor pad_rows_by_copy(const torch::Tensor& input, + int64_t padded_rows) { + CHECK_GE(padded_rows, input.size(0)); + if (padded_rows == input.size(0)) { + return input; + } + + auto output_shape = input.sizes().vec(); + output_shape[0] = padded_rows; + auto output = torch::empty(output_shape, input.options()); + output.slice(0, 0, input.size(0)).copy_(input); + output.slice(0, input.size(0), padded_rows).zero_(); + return output; } FlashComm1Context build_flash_comm1_context( @@ -178,7 +183,7 @@ FlashComm1Context build_flash_comm1_context( ctx.tp_world_size * kFc1LocalTokenAlignment; ctx.padded_num_tokens = round_up_to_multiple(num_tokens, token_alignment); ctx.pad_size = ctx.padded_num_tokens - num_tokens; - ctx.local_num_tokens = ctx.local_num_tokens_for_rank(ctx.tp_rank); + ctx.local_num_tokens = local_num_tokens_for_rank(ctx, ctx.tp_rank); ctx.padded_local_num_tokens = ctx.padded_num_tokens / ctx.tp_world_size; return ctx; @@ -198,7 +203,7 @@ FlashComm1Context build_flash_comm1_context(int32_t num_tokens, torch::Tensor shard_sequence(const torch::Tensor& input, const FlashComm1Context& ctx) { - if (!ctx.is_sequence_sharded()) { + if (!is_sequence_sharded(ctx)) { return input; } @@ -216,7 +221,7 @@ torch::Tensor shard_sequence(const torch::Tensor& input, torch::Tensor gather_sequence(const torch::Tensor& input, const FlashComm1Context& ctx) { - if (!ctx.is_sequence_sharded()) { + if (!is_sequence_sharded(ctx)) { return input; } @@ -224,10 +229,10 @@ torch::Tensor gather_sequence(const torch::Tensor& input, const int32_t expected_even_size = ctx.padded_num_tokens / ctx.tp_world_size; const int32_t expected_local_size = ctx.local_num_tokens; - std::vector token_num_list = ctx.token_num_list(); + std::vector token_nums = token_num_list(ctx); if (ctx.pad_size > 0 && current_local_size == expected_even_size) { for (int32_t i = 0; i < ctx.tp_world_size; ++i) { - token_num_list[i] = expected_even_size; + token_nums[i] = expected_even_size; } } else { CHECK_EQ(current_local_size, expected_local_size) @@ -236,7 +241,7 @@ torch::Tensor gather_sequence(const torch::Tensor& input, << ", rank=" << ctx.tp_rank << ", world=" << ctx.tp_world_size; } - auto gathered = parallel_state::gather(input, ctx.tp_group, token_num_list); + auto gathered = parallel_state::gather(input, ctx.tp_group, token_nums); if (ctx.pad_size > 0 && gathered.size(0) > ctx.original_num_tokens) { return gathered.slice(0, 0, ctx.original_num_tokens); @@ -271,7 +276,7 @@ namespace { torch::Tensor reduce_scatter_padded_local(const torch::Tensor& input, const FlashComm1Context& ctx) { CHECK(ctx.tp_group); - CHECK(ctx.is_sequence_sharded()); + CHECK(is_sequence_sharded(ctx)); CHECK_EQ(input.size(0), ctx.original_num_tokens) << "FC1 row-parallel reduce_scatter expects full real-token output " << "before communication."; @@ -307,7 +312,7 @@ torch::Tensor maybe_pad_and_reduce(torch::Tensor input, mode == RowParallelReduceMode::MATMUL_REDUCE_SCATTER) << "Unsupported row-parallel reduce mode."; - if (!ctx.is_sequence_sharded()) { + if (!is_sequence_sharded(ctx)) { if (ctx.tp_group && ctx.tp_group->world_size() > 1) { return parallel_state::reduce(input, ctx.tp_group); } @@ -334,15 +339,15 @@ torch::Tensor maybe_chunk_residual(const torch::Tensor& residual, CHECK_LT(tp_rank, tp_world_size); const int32_t num_tokens = static_cast(residual.size(0)); const int64_t start = - shard_start_for_rank(num_tokens, tp_world_size, tp_rank); - const int64_t end = - start + local_num_tokens_for_rank(num_tokens, tp_world_size, tp_rank); + shard_start_for_rank_impl(num_tokens, tp_world_size, tp_rank); + const int64_t end = start + local_num_tokens_for_rank_impl( + num_tokens, tp_world_size, tp_rank); return residual.slice(0, start, end).contiguous(); } torch::Tensor maybe_shard_residual(const torch::Tensor& residual, const FlashComm1Context& ctx) { - if (!ctx.is_sequence_sharded()) { + if (!is_sequence_sharded(ctx)) { return residual; } if (residual.size(0) == ctx.padded_local_num_tokens) { diff --git a/xllm/core/common/flash_comm1_context.h b/xllm/core/common/flash_comm1_context.h index bb0ccf4435..bf39e5d6cb 100644 --- a/xllm/core/common/flash_comm1_context.h +++ b/xllm/core/common/flash_comm1_context.h @@ -28,42 +28,32 @@ struct AttentionMetadata; } enum class RowParallelReduceMode : int8_t { - NONE = 0, - ALL_REDUCE = 1, - REDUCE_SCATTER = 2, - MATMUL_REDUCE_SCATTER = 3, + NONE = 0, + ALL_REDUCE = 1, + REDUCE_SCATTER = 2, + MATMUL_REDUCE_SCATTER = 3, }; struct FlashComm1Context { - bool enabled = false; - int32_t tp_rank = 0; - int32_t tp_world_size = 1; - int32_t original_num_tokens = 0; - int32_t padded_num_tokens = 0; - int32_t local_num_tokens = 0; - int32_t padded_local_num_tokens = 0; - int32_t pad_size = 0; - bool is_prefill = true; - bool enable_mmrs_fusion = false; - std::string mmrs_comm_mode = "aiv"; - ProcessGroup* tp_group = nullptr; - - FlashComm1Context() = default; - - bool is_sequence_sharded() const { return enabled && tp_world_size > 1; } - - int32_t local_num_tokens_for_rank(int32_t rank) const; - std::vector token_num_list() const; - - int64_t get_shard_start() const; - int64_t get_shard_end() const; + bool enabled = false; + int32_t tp_rank = 0; + int32_t tp_world_size = 1; + int32_t original_num_tokens = 0; + int32_t padded_num_tokens = 0; + int32_t local_num_tokens = 0; + int32_t padded_local_num_tokens = 0; + int32_t pad_size = 0; + bool is_prefill = true; + bool enable_mmrs_fusion = false; + std::string mmrs_comm_mode = "aiv"; + ProcessGroup* tp_group = nullptr; }; struct FlashComm1Options { - bool enable_flashcomm1 = false; - int32_t min_prefill_tokens = 8192; - bool enable_mmrs_fusion = false; - std::string mmrs_comm_mode = "aiv"; + bool enable_flashcomm1 = false; + int32_t min_prefill_tokens = 8192; + bool enable_mmrs_fusion = false; + std::string mmrs_comm_mode = "aiv"; }; class FlashComm1ContextScope { @@ -80,6 +70,19 @@ class FlashComm1ContextScope { const FlashComm1Context* get_current_flash_comm1_context(); +bool is_sequence_sharded(const FlashComm1Context& ctx); + +int32_t local_num_tokens_for_rank(const FlashComm1Context& ctx, int32_t rank); + +std::vector token_num_list(const FlashComm1Context& ctx); + +int64_t get_shard_start(const FlashComm1Context& ctx); + +int64_t get_shard_end(const FlashComm1Context& ctx); + +torch::Tensor pad_rows_by_copy(const torch::Tensor& input, + int64_t padded_rows); + FlashComm1Context build_flash_comm1_context( int32_t num_tokens, bool is_prefill, @@ -91,24 +94,28 @@ FlashComm1Context build_flash_comm1_context( const ParallelArgs& parallel_args, const FlashComm1Options& options); -torch::Tensor shard_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); +torch::Tensor shard_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx); -torch::Tensor gather_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); +torch::Tensor gather_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx); -torch::Tensor gather_and_unpad_sequence(const torch::Tensor& input, const FlashComm1Context& ctx); +torch::Tensor gather_and_unpad_sequence(const torch::Tensor& input, + const FlashComm1Context& ctx); -torch::Tensor maybe_pad_for_reduce(const torch::Tensor& input, const FlashComm1Context& ctx); +torch::Tensor maybe_pad_for_reduce(const torch::Tensor& input, + const FlashComm1Context& ctx); torch::Tensor maybe_pad_and_reduce(torch::Tensor input, - const FlashComm1Context& ctx, - RowParallelReduceMode mode); + const FlashComm1Context& ctx, + RowParallelReduceMode mode); RowParallelReduceMode row_parallel_reduce_mode_for_fc1( const FlashComm1Context& ctx); torch::Tensor maybe_chunk_residual(const torch::Tensor& residual, - int32_t tp_rank, - int32_t tp_world_size); + int32_t tp_rank, + int32_t tp_world_size); torch::Tensor maybe_shard_residual(const torch::Tensor& residual, const FlashComm1Context& ctx); diff --git a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp index 938ac1bf73..f809f5dfc1 100644 --- a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp +++ b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp @@ -39,6 +39,12 @@ bool can_call_torch_npu_mmrs( << "FC1 MMRS torch_npu skipped: output tensor is missing."; return false; } + if (!a.defined() || !b.defined()) { + LOG_FIRST_N(WARNING, 8) + << "FC1 MMRS torch_npu skipped: input tensor is missing. a_defined=" + << a.defined() << ", b_defined=" << b.defined(); + return false; + } if (a.dim() != 2 || b.dim() != 2 || output->dim() != 2) { LOG_FIRST_N(WARNING, 8) << "FC1 MMRS torch_npu skipped: expected 2D tensors, got a_dim=" diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index 760ba4db0b..8ea67acc40 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -97,14 +97,14 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states) { const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { h = gather_sequence(hidden_states, *fc1_ctx); } auto gate_up = gate_up_proj_->forward(h); if (is_smoothquant_) { - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite DenseMLP.down_proj(smoothquant): input=" << gate_up.sizes(); @@ -126,7 +126,7 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states) { act_->forward(gate_up, output); - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite DenseMLP.down_proj: input=" << output.sizes(); return down_proj_->forward( diff --git a/xllm/core/layers/common/linear.cpp b/xllm/core/layers/common/linear.cpp index 7f2420889b..e467ddb26f 100644 --- a/xllm/core/layers/common/linear.cpp +++ b/xllm/core/layers/common/linear.cpp @@ -21,6 +21,7 @@ limitations under the License. #include #include +#include "common/flash_comm1_context.h" #include "framework/parallel_state/parallel_args.h" #include "framework/parallel_state/parallel_state.h" #include "kernels/ops_api.h" @@ -61,21 +62,6 @@ inline bool is_unfused_checkpoint(const std::vector& scales) { scales.back() > std::numeric_limits::lowest(); } -torch::Tensor pad_rows_by_copy(const torch::Tensor& input, - int64_t padded_rows) { - CHECK_GE(padded_rows, input.size(0)); - if (padded_rows == input.size(0)) { - return input; - } - - auto output_shape = input.sizes().vec(); - output_shape[0] = padded_rows; - auto output = torch::empty(output_shape, input.options()); - output.slice(0, 0, input.size(0)).copy_(input); - output.slice(0, input.size(0), padded_rows).zero_(); - return output; -} - // Realigns FP8 partitions to a unified global scale to enable fusion. // Logic: // 1. Recover original values (FP8 -> FP16) using partition-specific scales. @@ -276,7 +262,7 @@ void log_mmrs_quant_skip(RowParallelReduceMode reduce_mode, << " path: fused matmul_reduce_scatter is currently wired only for " "non-quant linear. input=" << input.sizes() << ", sequence_sharded=" - << (fc1_ctx != nullptr && fc1_ctx->is_sequence_sharded()) + << (fc1_ctx != nullptr && is_sequence_sharded(*fc1_ctx)) << ", enable_mmrs_fusion=" << (fc1_ctx != nullptr && fc1_ctx->enable_mmrs_fusion); } @@ -1577,7 +1563,7 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, ? std::optional(bias_) : std::nullopt; - const bool skip_scatter = fc1_ctx && fc1_ctx->is_sequence_sharded(); + const bool skip_scatter = fc1_ctx && is_sequence_sharded(*fc1_ctx); torch::Tensor output; if (quant_args_.quant_method() == kQuantMethodSmoothquant) { @@ -1685,7 +1671,8 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, if (!input_is_parallelized_ && !skip_scatter) { input = xllm::parallel_state::scatter(input, process_group_); } - if (wants_mmrs(reduce_mode) && fc1_ctx && fc1_ctx->is_sequence_sharded() && + if (wants_mmrs(reduce_mode) && fc1_ctx && + is_sequence_sharded(*fc1_ctx) && fc1_ctx->enable_mmrs_fusion) { bool can_try_mmrs = input.dim() == 2 && input.size(0) == fc1_ctx->original_num_tokens && @@ -1751,7 +1738,7 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, << "FC1 MMRS skipped before row-parallel matmul: fc1_ctx=" << (fc1_ctx != nullptr) << ", sequence_sharded=" - << (fc1_ctx != nullptr && fc1_ctx->is_sequence_sharded()) + << (fc1_ctx != nullptr && is_sequence_sharded(*fc1_ctx)) << ", enable_mmrs_fusion=" << (fc1_ctx != nullptr && fc1_ctx->enable_mmrs_fusion) << ", reduce_mode=" << static_cast(reduce_mode) diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index e969b5a62e..fdf2e603f5 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -535,7 +535,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const ModelInputParams& input_params) { const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { h = gather_sequence(hidden_states, *fc1_ctx); } @@ -914,7 +914,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( rearranged_norm = rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); } - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite Qwen3GatedDeltaNet.o_proj: input=" << rearranged_norm.sizes(); diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index 646c4d564d..7fe0eb020b 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -150,7 +150,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); torch::Tensor h = hidden_states; - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { h = gather_sequence(hidden_states, *fc1_ctx); } @@ -179,7 +179,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( attn_->forward(attn_metadata, q_flat, k_flat, v_flat, kv_cache)); out = out * torch::sigmoid(gate.view({T, q_size_})); - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite Qwen3NextAttention.o_proj(mrope): input=" << out.sizes(); @@ -216,7 +216,7 @@ torch::Tensor Qwen3NextAttentionImpl::forward( out = out * torch::sigmoid(gate); } - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { LOG_FIRST_N(INFO, 16) << "FC1 MMRS callsite Qwen3NextAttention.o_proj: input=" << out.sizes(); diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp index 81426acee9..0eaf5e4969 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.cpp @@ -121,11 +121,11 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( residual = x; x = std::get<0>(input_norm_->forward(x)); } else { - if (fc1_ctx && fc1_ctx->is_sequence_sharded() && + if (fc1_ctx && is_sequence_sharded(*fc1_ctx) && residual.value().size(0) != x.size(0)) { residual = maybe_shard_residual(residual.value(), *fc1_ctx); } - if (fc1_ctx && fc1_ctx->is_sequence_sharded()) { + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { CHECK_EQ(residual.value().size(0), x.size(0)) << "FC1 input residual and hidden states must share the same " << "padded local sequence layout."; @@ -142,7 +142,7 @@ torch::Tensor Qwen3HybridDecoderLayerImplBase::forward( } // Before post_norm, ensure residual shape matches x shape - if (fc1_ctx && fc1_ctx->is_sequence_sharded() && residual.has_value() && + if (fc1_ctx && is_sequence_sharded(*fc1_ctx) && residual.has_value() && residual.value().size(0) != x.size(0)) { residual = maybe_shard_residual(residual.value(), *fc1_ctx); CHECK_EQ(residual.value().size(0), x.size(0)) diff --git a/xllm/models/llm/qwen3_next_hybrid_base.h b/xllm/models/llm/qwen3_next_hybrid_base.h index 0ba313813e..f1ee546e94 100644 --- a/xllm/models/llm/qwen3_next_hybrid_base.h +++ b/xllm/models/llm/qwen3_next_hybrid_base.h @@ -116,7 +116,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { h = embed_tokens_(tokens); } - if (fc1_ctx.is_sequence_sharded()) { + if (is_sequence_sharded(fc1_ctx)) { h = shard_sequence(h, fc1_ctx); } @@ -146,7 +146,7 @@ class Qwen3HybridModelImplBase : public Qwen3HybridModelModule { } auto [hidden_states, residual_out] = norm_->forward(h, residual); h = hidden_states; - if (fc1_ctx.is_sequence_sharded()) { + if (is_sequence_sharded(fc1_ctx)) { h = gather_and_unpad_sequence(h, fc1_ctx); } return ModelOutput(h); From f9b49341440c91ec9cafbbe213f21a22be3cde02 Mon Sep 17 00:00:00 2001 From: Cooofish <2716856597@qq.com> Date: Tue, 7 Jul 2026 15:43:24 +0800 Subject: [PATCH 5/6] bugfix: drop cmake and third_party changes from flashcomm1 pr. - Restore cmake helper files to match upstream main. - Restore third_party submodule pointers to match upstream main so the PR only carries FC1 source changes. --- cmake/FindRust.cmake | 10 ++++++---- cmake/cargo_library.cmake | 4 +--- cmake/cargo_shared_library.cmake | 4 +--- third_party/xllm_atb_layers | 2 +- third_party/xllm_ops | 2 +- 5 files changed, 10 insertions(+), 12 deletions(-) diff --git a/cmake/FindRust.cmake b/cmake/FindRust.cmake index b3e0f995fb..8c771250dd 100644 --- a/cmake/FindRust.cmake +++ b/cmake/FindRust.cmake @@ -10,10 +10,12 @@ else() set(USER_HOME "$ENV{HOME}") endif() -if(NOT "$ENV{CARGO_HOME}" STREQUAL "") - set(CARGO_HOME "$ENV{CARGO_HOME}" CACHE PATH "Rust Cargo Home" FORCE) -elseif(NOT DEFINED CARGO_HOME) - set(CARGO_HOME "${USER_HOME}/.cargo") +if(NOT DEFINED CARGO_HOME) + if("$ENV{CARGO_HOME}" STREQUAL "") + set(CARGO_HOME "${USER_HOME}/.cargo") + else() + set(CARGO_HOME "$ENV{CARGO_HOME}") + endif() endif() # Find cargo executable diff --git a/cmake/cargo_library.cmake b/cmake/cargo_library.cmake index fadc5725ce..bf34819525 100644 --- a/cmake/cargo_library.cmake +++ b/cmake/cargo_library.cmake @@ -65,9 +65,7 @@ function(cargo_library) file(GLOB_RECURSE LIB_SOURCES "*.rs") - set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env - "CARGO_HOME=${CARGO_HOME}" - "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") + set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") # build the library target with cargo set(STATIC_LIB_NAME diff --git a/cmake/cargo_shared_library.cmake b/cmake/cargo_shared_library.cmake index d92e7f7ab1..8bac49f016 100644 --- a/cmake/cargo_shared_library.cmake +++ b/cmake/cargo_shared_library.cmake @@ -65,9 +65,7 @@ function(cargo_shared_library) file(GLOB_RECURSE LIB_SOURCES "*.rs") - set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env - "CARGO_HOME=${CARGO_HOME}" - "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") + set(CARGO_ENV_COMMAND ${CMAKE_COMMAND} -E env "CARGO_TARGET_DIR=${CMAKE_CURRENT_BINARY_DIR}") # build the library target with cargo set(SHARED_LIB_NAME diff --git a/third_party/xllm_atb_layers b/third_party/xllm_atb_layers index 5b50a7e167..4b926ee633 160000 --- a/third_party/xllm_atb_layers +++ b/third_party/xllm_atb_layers @@ -1 +1 @@ -Subproject commit 5b50a7e167267d4482c665717f25e1ef066a8e86 +Subproject commit 4b926ee633ed11677890d7e2ddbad7eb6dd249f2 diff --git a/third_party/xllm_ops b/third_party/xllm_ops index e7d82584a7..bbbdfe3818 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit e7d82584a7c3043b4b5a08f60643e04537771147 +Subproject commit bbbdfe3818149b6d4d725e1fbf44149c4c01a1a6 From bf3d04fac4fb55445e615143ff2afcce8576ff10 Mon Sep 17 00:00:00 2001 From: Cooofish <2716856597@qq.com> Date: Wed, 8 Jul 2026 09:59:04 +0800 Subject: [PATCH 6/6] refactor: reduce flashcomm1 pr diff scope Remove non-essential build/runtime/header changes from the FlashComm1 PR diff, drop unused MMRS parameters, and remove noisy success-path FC1 logs. Keep the FC1 and MMRS functional paths intact while reducing unrelated review surface. --- CMakeLists.txt | 46 +------------------ xllm/core/common/flash_comm1_context.cpp | 5 -- xllm/core/kernels/npu/matmul.cpp | 1 + .../kernels/npu/matmul_reduce_scatter.cpp | 11 ----- xllm/core/kernels/param.h | 7 --- xllm/core/layers/common/dense_mlp.cpp | 5 -- xllm/core/layers/common/dense_mlp.h | 1 - xllm/core/layers/common/linear.cpp | 5 -- .../npu_torch/qwen3_gated_delta_net_base.cpp | 3 -- .../npu_torch/qwen3_gated_delta_net_base.h | 1 - .../layers/npu_torch/qwen3_next_attention.cpp | 6 --- .../layers/npu_torch/qwen3_next_attention.h | 9 ++-- .../qwen3_next_hybrid_decoder_layer_base.h | 25 +++++----- xllm/core/runtime/llm_worker_impl.cpp | 6 +-- 14 files changed, 20 insertions(+), 111 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index f4cfa3cd0b..9090501db7 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -334,53 +334,9 @@ find_package(benchmark CONFIG REQUIRED) find_package(nlohmann_json CONFIG REQUIRED) find_package(OpenCV CONFIG REQUIRED) find_package(FFMPEG REQUIRED) -find_package(Python COMPONENTS Interpreter Development REQUIRED) +find_package(Python COMPONENTS Development REQUIRED) find_package(pybind11 CONFIG REQUIRED) -if(USE_NPU) - function(xllm_set_env_from_python VAR PYTHON_CODE) - if("$ENV{${VAR}}" STREQUAL "") - execute_process( - COMMAND ${Python_EXECUTABLE} -c "${PYTHON_CODE}" - OUTPUT_VARIABLE XLLM_INFERRED_ENV_VALUE - OUTPUT_STRIP_TRAILING_WHITESPACE - ERROR_VARIABLE XLLM_INFERRED_ENV_ERROR - RESULT_VARIABLE XLLM_INFERRED_ENV_RESULT - ) - if(NOT XLLM_INFERRED_ENV_RESULT EQUAL 0 OR "${XLLM_INFERRED_ENV_VALUE}" STREQUAL "") - message(FATAL_ERROR "Failed to infer ${VAR}: ${XLLM_INFERRED_ENV_ERROR}") - endif() - set(ENV{${VAR}} "${XLLM_INFERRED_ENV_VALUE}") - message(STATUS "Inferred ${VAR}: $ENV{${VAR}}") - endif() - endfunction() - - if("$ENV{PYTHON_INCLUDE_PATH}" STREQUAL "") - set(ENV{PYTHON_INCLUDE_PATH} "${Python_INCLUDE_DIRS}") - endif() - if("$ENV{PYTHON_LIB_PATH}" STREQUAL "") - execute_process( - COMMAND ${Python_EXECUTABLE} -c "import sysconfig; print(sysconfig.get_config_var('LIBDIR') or '')" - OUTPUT_VARIABLE XLLM_PYTHON_LIB_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE - ) - set(ENV{PYTHON_LIB_PATH} "${XLLM_PYTHON_LIB_PATH}") - endif() - xllm_set_env_from_python(PYTORCH_INSTALL_PATH "import importlib.util; spec = importlib.util.find_spec('torch'); print(spec.submodule_search_locations[0])") - xllm_set_env_from_python(PYTORCH_NPU_INSTALL_PATH "import importlib.util; spec = importlib.util.find_spec('torch_npu'); print(spec.submodule_search_locations[0])") - if(NOT EXISTS "$ENV{PYTORCH_NPU_INSTALL_PATH}/include/torch_npu/torch_npu.h" - AND EXISTS "/usr/local/libtorch_npu/include/torch_npu/torch_npu.h") - set(ENV{PYTORCH_NPU_INSTALL_PATH} "/usr/local/libtorch_npu") - message(STATUS "Using PYTORCH_NPU_INSTALL_PATH with compatibility header: $ENV{PYTORCH_NPU_INSTALL_PATH}") - endif() - if("$ENV{LIBTORCH_ROOT}" STREQUAL "") - set(ENV{LIBTORCH_ROOT} "$ENV{PYTORCH_INSTALL_PATH}") - endif() - if("$ENV{NPU_HOME_PATH}" STREQUAL "" AND NOT "$ENV{ASCEND_HOME_PATH}" STREQUAL "") - set(ENV{NPU_HOME_PATH} "$ENV{ASCEND_HOME_PATH}") - endif() -endif() - if (USE_CXX11_ABI) # only use jemalloc if using the new C++-11 ABI find_package(Jemalloc) diff --git a/xllm/core/common/flash_comm1_context.cpp b/xllm/core/common/flash_comm1_context.cpp index a11fd01c20..f54489dc6c 100644 --- a/xllm/core/common/flash_comm1_context.cpp +++ b/xllm/core/common/flash_comm1_context.cpp @@ -289,11 +289,6 @@ torch::Tensor reduce_scatter_padded_local(const torch::Tensor& input, auto output_shape = padded_input.sizes().vec(); output_shape[0] = ctx.padded_local_num_tokens; torch::Tensor output = torch::empty(output_shape, padded_input.options()); - LOG_FIRST_N(INFO, 16) - << "FC1 reduce_scatter active: rank=" << ctx.tp_rank - << ", original_tokens=" << ctx.original_num_tokens - << ", padded_tokens=" << ctx.padded_num_tokens - << ", local_tokens=" << ctx.padded_local_num_tokens; ctx.tp_group->reduce_scatter(padded_input, output); return output; } diff --git a/xllm/core/kernels/npu/matmul.cpp b/xllm/core/kernels/npu/matmul.cpp index eee3d9f13e..71d3408959 100644 --- a/xllm/core/kernels/npu/matmul.cpp +++ b/xllm/core/kernels/npu/matmul.cpp @@ -14,6 +14,7 @@ limitations under the License. ==============================================================================*/ #include "npu_ops_api.h" +#include "ops_npu/npu_ops.h" namespace xllm::kernel::npu { diff --git a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp index f809f5dfc1..c01c3c1e82 100644 --- a/xllm/core/kernels/npu/matmul_reduce_scatter.cpp +++ b/xllm/core/kernels/npu/matmul_reduce_scatter.cpp @@ -147,17 +147,6 @@ torch::Tensor matmul_reduce_scatter( << ", dtype=" << a.scalar_type(); } - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS torch_npu hit: a=" << a.sizes() << ", b=" << b.sizes() - << ", expected_out=" << output->sizes() << ", dtype=" << a.scalar_type() - << ", rank=" << process_group->rank() - << ", world_size=" << process_group->world_size() - << ", reduce_op=" << (reduce_op.empty() ? "sum" : reduce_op) - << ", comm_turn=" << comm_turn << ", stream_mode=" << stream_mode - << ", comm_mode=" << (comm_mode.empty() ? "none" : comm_mode) - << ", effective_comm_mode=" - << (effective_comm_mode.empty() ? "none" : effective_comm_mode); - std::optional torch_comm_mode = std::nullopt; if (effective_comm_mode == "ai_cpu" || effective_comm_mode == "aiv") { torch_comm_mode = c10::string_view(effective_comm_mode); diff --git a/xllm/core/kernels/param.h b/xllm/core/kernels/param.h index 028e84cd5b..c1693d1615 100644 --- a/xllm/core/kernels/param.h +++ b/xllm/core/kernels/param.h @@ -317,18 +317,11 @@ struct MatmulParams { }; struct MatmulReduceScatterParams { - // RFC-visible inputs. For the currently wired BF16/FP16 path, b is passed in - // K-N layout to match torch_npu::npu_mm_reduce_scatter_base. torch::Tensor a; torch::Tensor b; std::optional bias; ProcessGroup* process_group = nullptr; - int64_t original_num_tokens = 0; - std::optional deq_scale; - std::optional output_dtype; - // Optional compatibility fields used by callers for expected shape checks. - // The torch_npu op allocates and returns the output tensor itself. std::optional output; std::string reduce_op = "sum"; int64_t comm_turn = 0; diff --git a/xllm/core/layers/common/dense_mlp.cpp b/xllm/core/layers/common/dense_mlp.cpp index 8ea67acc40..92014373eb 100644 --- a/xllm/core/layers/common/dense_mlp.cpp +++ b/xllm/core/layers/common/dense_mlp.cpp @@ -105,9 +105,6 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states) { if (is_smoothquant_) { if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS callsite DenseMLP.down_proj(smoothquant): input=" - << gate_up.sizes(); return down_proj_->forward( gate_up, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } @@ -127,8 +124,6 @@ torch::Tensor DenseMLPImpl::forward(torch::Tensor hidden_states) { act_->forward(gate_up, output); if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS callsite DenseMLP.down_proj: input=" << output.sizes(); return down_proj_->forward( output, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } diff --git a/xllm/core/layers/common/dense_mlp.h b/xllm/core/layers/common/dense_mlp.h index 291f0d257a..7aa8386111 100644 --- a/xllm/core/layers/common/dense_mlp.h +++ b/xllm/core/layers/common/dense_mlp.h @@ -20,7 +20,6 @@ limitations under the License. #include #include "activation.h" -#include "common/flash_comm1_context.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" #include "framework/quant_args.h" diff --git a/xllm/core/layers/common/linear.cpp b/xllm/core/layers/common/linear.cpp index e467ddb26f..659bb6c823 100644 --- a/xllm/core/layers/common/linear.cpp +++ b/xllm/core/layers/common/linear.cpp @@ -1693,15 +1693,10 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input, mmrs_params.b = mmrs_weight_transposed(); mmrs_params.bias = bias; mmrs_params.process_group = process_group_; - mmrs_params.original_num_tokens = fc1_ctx->original_num_tokens; mmrs_params.output = mmrs_output; mmrs_params.comm_mode = fc1_ctx->mmrs_comm_mode; output = xllm::kernel::matmul_reduce_scatter(mmrs_params); if (output.sizes() == torch::IntArrayRef(output_shape)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS row-parallel fused output accepted: input=" - << input.sizes() << ", weight=" << weight_.sizes() - << ", output=" << output.sizes(); return output; } LOG_FIRST_N(WARNING, 8) diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index fdf2e603f5..51a631f47e 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -915,9 +915,6 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); } if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS callsite Qwen3GatedDeltaNet.o_proj: input=" - << rearranged_norm.sizes(); return o_proj_->forward( rearranged_norm, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h index 9dbf0bcfaf..afbda14b77 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.h @@ -22,7 +22,6 @@ limitations under the License. #include #include "attention.h" -#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp index 7fe0eb020b..22bf7f7817 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.cpp +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.cpp @@ -180,9 +180,6 @@ torch::Tensor Qwen3NextAttentionImpl::forward( out = out * torch::sigmoid(gate.view({T, q_size_})); if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS callsite Qwen3NextAttention.o_proj(mrope): input=" - << out.sizes(); return o_proj_->forward( out, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } @@ -217,9 +214,6 @@ torch::Tensor Qwen3NextAttentionImpl::forward( } if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { - LOG_FIRST_N(INFO, 16) - << "FC1 MMRS callsite Qwen3NextAttention.o_proj: input=" - << out.sizes(); return o_proj_->forward( out, row_parallel_reduce_mode_for_fc1(*fc1_ctx)); } diff --git a/xllm/core/layers/npu_torch/qwen3_next_attention.h b/xllm/core/layers/npu_torch/qwen3_next_attention.h index f6c3ff1a5a..45347fb919 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_attention.h +++ b/xllm/core/layers/npu_torch/qwen3_next_attention.h @@ -20,7 +20,6 @@ limitations under the License. #include #include "attention.h" -#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_args.h" #include "framework/parallel_state/parallel_args.h" @@ -44,10 +43,10 @@ class Qwen3NextAttentionImpl : public torch::nn::Module { int32_t layer_id); torch::Tensor forward(const torch::Tensor& positions, - const torch::Tensor& hidden_states, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const torch::Tensor& mrope_cos_sin); + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const torch::Tensor& mrope_cos_sin); torch::Tensor build_mrope_cos_sin(const torch::Tensor& positions) const; diff --git a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h index e431a1698a..fb6d3a6c83 100644 --- a/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h +++ b/xllm/core/layers/npu_torch/qwen3_next_hybrid_decoder_layer_base.h @@ -20,7 +20,6 @@ limitations under the License. #include #include -#include "common/flash_comm1_context.h" #include "framework/kv_cache/kv_cache.h" #include "framework/model/model_input_params.h" #include "framework/model_context.h" @@ -39,12 +38,12 @@ class Qwen3HybridDecoderLayerModule : public torch::nn::Module { virtual void load_state_dict(const StateDict& state_dict) = 0; virtual void verify_loaded_weights(const std::string& prefix) const = 0; virtual torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) = 0; + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) = 0; virtual torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const { return {}; @@ -66,12 +65,12 @@ class Qwen3HybridDecoderLayerImplBase : public Qwen3HybridDecoderLayerModule { void verify_loaded_weights(const std::string& prefix) const override; torch::Tensor forward(torch::Tensor& x, - std::optional& residual, - torch::Tensor& positions, - const AttentionMetadata& attn_metadata, - KVCache& kv_cache, - const ModelInputParams& input_params, - const torch::Tensor& mrope_cos_sin = {}) override; + std::optional& residual, + torch::Tensor& positions, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params, + const torch::Tensor& mrope_cos_sin = {}) override; torch::Tensor build_mrope_cos_sin( const torch::Tensor& positions) const override; diff --git a/xllm/core/runtime/llm_worker_impl.cpp b/xllm/core/runtime/llm_worker_impl.cpp index 439da3aa55..96c6b4b2d0 100644 --- a/xllm/core/runtime/llm_worker_impl.cpp +++ b/xllm/core/runtime/llm_worker_impl.cpp @@ -250,11 +250,9 @@ std::optional LLMWorkerImpl::step_internal( eplb_executor_->eplb_execute(input.input_params.expert.eplb_info); } - ModelInputParams model_input_params = input.input_params; - // call model executor forward to get hidden states auto model_output = model_executor_->forward( - input.token_ids, input.positions, kv_caches_, model_input_params); + input.token_ids, input.positions, kv_caches_, input.input_params); if (!model_output.hidden_states.defined()) { return std::nullopt; } @@ -263,7 +261,7 @@ std::optional LLMWorkerImpl::step_internal( torch::Tensor selected_hidden_from_lm_head; if (sampling_params.selected_token_idxes.defined()) { torch::Tensor selected_token_idxes = sampling_params.selected_token_idxes; - if (model_output.hidden_states.defined() && selected_token_idxes.defined() && + if (model_output.hidden_states.defined() && selected_token_idxes.device() != model_output.hidden_states.device()) { selected_token_idxes = selected_token_idxes .to(model_output.hidden_states.device(),