feat: add FlashComm1 and MMRS fusion support for Qwen3.5 NPU.

docs/src/content/docs/en/features/flashcomm.md

@@ -45,6 +45,7 @@ FlashComm is usually not suitable, or has limited benefit, for:
 - Long-output workloads where decode dominates total latency.
 - `TP < 8`, `dp > 1`, or `cp > 1`, which are not enabled by the default activation rule.
 - Quantized row-parallel paths. MMRS is currently wired only for the normal BF16/FP16 matmul path.
+- MoE models. FC1 currently supports Dense models only and is automatically disabled when routed experts are present.
 
 ## Usage
 
@@ -90,6 +91,8 @@ In most cases, keep `mmrs_comm_mode=aiv`. If some shapes hit AICore errors on th
 --mmrs_comm_mode=ai_cpu
 ```
 
+When using `mmrs_comm_mode=aiv`, load an ops-transformers 9.1.0 or newer operator library that includes the MMRS AIV fix.
+
 ## Recommended Configuration
 
 Use the following table as the initial tuning guidance:

docs/src/content/docs/zh/features/flashcomm.md

@@ -45,6 +45,7 @@ FlashComm 不适合或收益有限的场景：
 - 高 decode 占比场景，例如长输出请求。整体 latency 可能主要由 decode 决定。
 - `TP < 8`、`dp > 1` 或 `cp > 1` 的场景，当前默认不会启用。
 - 量化 row-parallel 路径。当前 MMRS 只接入普通 BF16/FP16 matmul 路径。
+- MoE 模型。当前 FC1 仅支持 Dense 模型；检测到 routed experts 时会自动关闭 FC1。
 
 ## 使用方式
 
@@ -90,6 +91,8 @@ FlashComm 和 MMRS 融合算子都默认关闭。推荐在长 prefill、TP=8 或
 --mmrs_comm_mode=ai_cpu
 ```
 
+使用 `mmrs_comm_mode=aiv` 时，需要加载包含 MMRS AIV 修复的 ops-transformers 9.1.0 或更高版本算子库。
+
 ## 最优配置建议
 
 推荐从以下配置开始评估：

xllm/core/common/CMakeLists.txt

@@ -17,16 +17,19 @@ cc_library(
     types.h
     device_monitor.h
     version_singleton.h
+    flash_comm1_context.h
   SRCS
     etcd_client.cpp
     metrics.cpp
     $<$<BOOL:${USE_NPU}>:mspti_helper.cpp>
     options.cpp
     rate_limiter.cpp
     device_monitor.cpp
+    flash_comm1_context.cpp
   DEPS
     :config
     util
+    parallel_state
     absl::random_random
     absl::strings
     torch

xllm/core/common/flash_comm1_context.cpp

@@ -0,0 +1,245 @@
+/* Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "flash_comm1_context.h"
+
+#include <glog/logging.h>
+
+#include <algorithm>
+
+#include "framework/parallel_state/parallel_state.h"
+
+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);
+  const int32_t remainder = value % multiple;
+  return remainder == 0 ? value : value + multiple - remainder;
+}
+
+}  // 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;
+}
+
+bool is_sequence_sharded(const FlashComm1Context& ctx) {
+  return ctx.enabled && ctx.tp_world_size > 1;
+}
+
+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(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)
+  return ctx;
+#endif
+
+  if (parallel_args.dp_size() != 1 || parallel_args.cp_size() != 1) {
+    return ctx;
+  }
+
+  if (actual_tp_size < kFc1MinTpSize) {
+    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.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.padded_local_num_tokens = ctx.padded_num_tokens / ctx.tp_world_size;
+
+  return ctx;
+}
+
+torch::Tensor shard_sequence(const torch::Tensor& input,
+                             const FlashComm1Context& ctx) {
+  if (!is_sequence_sharded(ctx)) {
+    return input;
+  }
+
+  CHECK_EQ(input.size(0), ctx.original_num_tokens);
+  const int64_t shard_start =
+      static_cast<int64_t>(ctx.tp_rank) * ctx.padded_local_num_tokens;
+  const int64_t shard_end = shard_start + ctx.padded_local_num_tokens;
+  const int64_t valid_end =
+      std::min(shard_end, static_cast<int64_t>(ctx.original_num_tokens));
+
+  if (valid_end == shard_end) {
+    return input.slice(0, shard_start, shard_end).contiguous();
+  }
+
+  auto output_shape = input.sizes().vec();
+  output_shape[0] = ctx.padded_local_num_tokens;
+  torch::Tensor output = torch::zeros(output_shape, input.options());
+  if (valid_end > shard_start) {
+    output.slice(0, 0, valid_end - shard_start)
+        .copy_(input.slice(0, shard_start, valid_end));
+  }
+  return output;
+}
+
+torch::Tensor gather_sequence(const torch::Tensor& input,
+                              const FlashComm1Context& ctx) {
+  if (!is_sequence_sharded(ctx)) {
+    return input;
+  }
+
+  const int32_t expected_local_size = ctx.padded_local_num_tokens;
+  CHECK_EQ(input.size(0), expected_local_size)
+      << "FC1 gather expects a padded local shard of " << expected_local_size
+      << " rows, got " << input.size(0) << ", rank=" << ctx.tp_rank
+      << ", world=" << ctx.tp_world_size;
+
+  const std::vector<int32_t> token_nums(ctx.tp_world_size, expected_local_size);
+
+  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);
+  }
+  return gathered;
+}
+
+namespace {
+
+torch::Tensor reduce_scatter_padded_local(const torch::Tensor& input,
+                                          const FlashComm1Context& ctx) {
+  CHECK(ctx.tp_group);
+  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.";
+
+  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());
+  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 (!is_sequence_sharded(ctx)) {
+    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_shard_residual(const torch::Tensor& residual,
+                                   const FlashComm1Context& ctx) {
+  if (!is_sequence_sharded(ctx)) {
+    return residual;
+  }
+  const int64_t num_tokens = residual.size(0);
+  CHECK(num_tokens == ctx.original_num_tokens ||
+        num_tokens == ctx.padded_local_num_tokens)
+      << "FC1 residual layout must be either full real-token or padded local "
+      << "sequence shard.";
+  if (num_tokens == ctx.original_num_tokens) {
+    return shard_sequence(residual, ctx);
+  }
+  return residual;
+}
+
+}  // namespace xllm

xllm/core/common/flash_comm1_context.h

@@ -0,0 +1,92 @@
+/* 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 <torch/torch.h>
+
+#include <string>
+
+#include "framework/parallel_state/parallel_args.h"
+
+namespace xllm {
+
+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 padded_local_num_tokens = 0;
+  int32_t pad_size = 0;
+  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";
+};
+
+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();
+
+bool is_sequence_sharded(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,
+                                            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 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_shard_residual(const torch::Tensor& residual,
+                                   const FlashComm1Context& ctx);
+
+}  // namespace xllm