feat: add flashcomm1 feature for deepseek v4 on NPU

tests/core/framework/parallel_state/CMakeLists.txt

@@ -1,5 +1,19 @@
 include(cc_test)
 
+# Host-buildable: exercises the FlashComm1 thread_local context RAII and the
+# device-free shard_dim0_padded sharding math. No device / process group.
+cc_test(
+  NAME
+    flash_comm1_context_test
+  SRCS
+    flash_comm1_context_test.cpp
+  DEPS
+    parallel_state
+    GTest::gtest_main
+    torch
+    glog::glog
+)
+
 if(USE_NPU)
   cc_test(
     NAME

tests/core/framework/parallel_state/flash_comm1_context_test.cpp

@@ -0,0 +1,155 @@
+/* Copyright 2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <gtest/gtest.h>
+#include <torch/torch.h>
+
+#include "core/framework/parallel_state/flash_comm1_context.h"
+#include "core/framework/parallel_state/parallel_state.h"
+
+namespace xllm {
+namespace parallel_state {
+
+// --- FlashComm1Context / FlashComm1Guard ------------------------------------
+// These exercise the thread_local RAII scoping only and need no device.
+
+TEST(FlashComm1ContextTest, InactiveByDefault) {
+  EXPECT_FALSE(flash_comm1_active());
+  EXPECT_FALSE(current_flash_comm1_context().enabled);
+}
+
+TEST(FlashComm1ContextTest, GuardPublishesAndRestores) {
+  ASSERT_FALSE(flash_comm1_active());
+  {
+    FlashComm1Guard guard(/*enabled=*/true,
+                          /*num_tokens=*/17,
+                          /*tp_rank=*/2,
+                          /*tp_world_size=*/4,
+                          /*tp_group=*/nullptr);
+    EXPECT_TRUE(flash_comm1_active());
+    const auto& ctx = current_flash_comm1_context();
+    EXPECT_TRUE(ctx.enabled);
+    EXPECT_EQ(ctx.num_tokens, 17);
+    EXPECT_EQ(ctx.tp_rank, 2);
+    EXPECT_EQ(ctx.tp_world_size, 4);
+    EXPECT_EQ(ctx.tp_group, nullptr);
+  }
+  // Restored to inactive on scope exit.
+  EXPECT_FALSE(flash_comm1_active());
+  EXPECT_FALSE(current_flash_comm1_context().enabled);
+}
+
+TEST(FlashComm1ContextTest, GuardsNestAndRestoreOuter) {
+  FlashComm1Guard outer(/*enabled=*/true,
+                        /*num_tokens=*/8,
+                        /*tp_rank=*/0,
+                        /*tp_world_size=*/2,
+                        /*tp_group=*/nullptr);
+  EXPECT_EQ(current_flash_comm1_context().num_tokens, 8);
+  {
+    FlashComm1Guard inner(/*enabled=*/true,
+                          /*num_tokens=*/32,
+                          /*tp_rank=*/1,
+                          /*tp_world_size=*/2,
+                          /*tp_group=*/nullptr);
+    EXPECT_EQ(current_flash_comm1_context().num_tokens, 32);
+    EXPECT_EQ(current_flash_comm1_context().tp_rank, 1);
+  }
+  // Inner guard restores the outer context, not the inactive default.
+  EXPECT_TRUE(flash_comm1_active());
+  EXPECT_EQ(current_flash_comm1_context().num_tokens, 8);
+  EXPECT_EQ(current_flash_comm1_context().tp_rank, 0);
+}
+
+TEST(FlashComm1ContextTest, DisabledGuardStillScopes) {
+  FlashComm1Guard guard(/*enabled=*/false,
+                        /*num_tokens=*/5,
+                        /*tp_rank=*/0,
+                        /*tp_world_size=*/1,
+                        /*tp_group=*/nullptr);
+  // A disabled guard reports inactive but its fields are still published, so a
+  // consumer that checks enabled first sees a consistent (off) state.
+  EXPECT_FALSE(flash_comm1_active());
+  EXPECT_FALSE(current_flash_comm1_context().enabled);
+}
+
+// --- shard_dim0_padded ------------------------------------------------------
+// No process group required: the sharding math runs on plain CPU tensors.
+
+TEST(ShardDim0PaddedTest, WorldSizeOneIsNoOp) {
+  auto x = torch::arange(6).reshape({3, 2});
+  auto shard = shard_dim0_padded(x, /*rank=*/0, /*world_size=*/1);
+  EXPECT_TRUE(torch::equal(shard, x));
+}
+
+TEST(ShardDim0PaddedTest, DivisibleSplitsEvenly) {
+  // 8 tokens, hidden=2, world=4 -> each rank owns 2 rows, no padding.
+  auto x = torch::arange(16).reshape({8, 2}).to(torch::kFloat32);
+  for (int32_t rank = 0; rank < 4; ++rank) {
+    auto shard = shard_dim0_padded(x, rank, /*world_size=*/4);
+    ASSERT_EQ(shard.size(0), 2);
+    ASSERT_EQ(shard.size(1), 2);
+    EXPECT_TRUE(torch::equal(shard, x.slice(0, rank * 2, rank * 2 + 2)));
+  }
+}
+
+TEST(ShardDim0PaddedTest, NonDivisiblePadsOnTrailingShards) {
+  // 5 tokens, world=2 -> padded to 6, chunk=3. rank0 = rows [0,3),
+  // rank1 = rows [3,6) where row 5 is zero padding.
+  auto x = torch::arange(5).reshape({5, 1}).to(torch::kFloat32);
+  auto shard0 = shard_dim0_padded(x, /*rank=*/0, /*world_size=*/2);
+  auto shard1 = shard_dim0_padded(x, /*rank=*/1, /*world_size=*/2);
+  ASSERT_EQ(shard0.size(0), 3);
+  ASSERT_EQ(shard1.size(0), 3);
+
+  EXPECT_TRUE(torch::equal(shard0, x.slice(0, 0, 3)));
+  // shard1 rows: [3, 4, pad(0)].
+  EXPECT_FLOAT_EQ(shard1[0].item<float>(), 3.0f);
+  EXPECT_FLOAT_EQ(shard1[1].item<float>(), 4.0f);
+  EXPECT_FLOAT_EQ(shard1[2].item<float>(), 0.0f);
+}
+
+TEST(ShardDim0PaddedTest, ConcatenatedShardsReconstructPaddedFull) {
+  // Concatenating every rank's shard in order == the padded full tensor. This
+  // mirrors what all_gather_dim0_unpad reconstructs before unpadding.
+  auto x = torch::arange(7).reshape({7, 1}).to(torch::kFloat32);
+  const int32_t world = 4;  // padded to 8, chunk = 2.
+  std::vector<torch::Tensor> shards;
+  for (int32_t rank = 0; rank < world; ++rank) {
+    shards.push_back(shard_dim0_padded(x, rank, world));
+  }
+  auto full = torch::cat(shards, /*dim=*/0);
+  ASSERT_EQ(full.size(0), 8);
+  // First 7 rows equal the original; the trailing row is zero padding.
+  EXPECT_TRUE(torch::equal(full.slice(0, 0, 7), x));
+  EXPECT_FLOAT_EQ(full[7].item<float>(), 0.0f);
+}
+
+TEST(ShardDim0PaddedTest, PreservesTrailingDims) {
+  // 3 tokens, shape [3, 2, 4], world=2 -> padded to 4, chunk=2, trailing dims
+  // (2, 4) untouched.
+  auto x = torch::randn({3, 2, 4});
+  auto shard = shard_dim0_padded(x, /*rank=*/1, /*world_size=*/2);
+  ASSERT_EQ(shard.dim(), 3);
+  EXPECT_EQ(shard.size(0), 2);
+  EXPECT_EQ(shard.size(1), 2);
+  EXPECT_EQ(shard.size(2), 4);
+  // rank1 row 0 == original row 2; row 1 is zero padding.
+  EXPECT_TRUE(torch::equal(shard[0], x[2]));
+  EXPECT_TRUE(torch::equal(shard[1], torch::zeros({2, 4})));
+}
+
+}  // namespace parallel_state
+}  // namespace xllm

xllm/core/common/global_flags.h

@@ -134,6 +134,11 @@ DECLARE_int32(expert_parallel_degree);
 
 DECLARE_string(rank_tablefile);
 
+// --- FlashComm1 (sequence parallel) ---
+DECLARE_bool(enable_flashcomm1);
+
+DECLARE_int32(flashcomm1_sp_min_token_num);
+
 constexpr int32_t kGraphExecutorLogVerboseLevel = 50;
 
 DECLARE_bool(enable_graph);

xllm/core/framework/config/parallel_config.cpp

@@ -69,6 +69,26 @@ DEFINE_bool(
     "Whether to enable dp load balance, if true, sequences within a single "
     "dp batch will be shuffled.");
 
+DEFINE_bool(
+    enable_flashcomm1,
+    false,
+    "Whether to enable FlashComm1 (sequence parallel): keep the residual "
+    "stream token-sharded across the TP group, all-gather to full tokens at "
+    "sequence-parallel boundaries (e.g. attention input), and replace the tail "
+    "all-reduce of eligible TP RowParallel layers with a padded "
+    "reduce-scatter(dim0). Requires tp>1 and benefits high-token "
+    "(prefill / chunked-prefill) batches. Defaults to false.");
+
+DEFINE_int32(
+    flashcomm1_sp_min_token_num,
+    1000,
+    "Minimum token count (dim0) for FlashComm1 sequence parallelism to engage. "
+    "Below this, the fixed overhead of the shard / all-gather / reduce-scatter "
+    "collectives dominates and SP is a net loss, so FlashComm1 is disabled for "
+    "the forward. This keeps SP on for prefill (many tokens) and off for decode "
+    "(few tokens). Set to 0 to disable the token gate entirely. Only meaningful "
+    "with enable_flashcomm1.");
+
 namespace xllm {
 
 void ParallelConfig::from_flags() {
@@ -85,6 +105,8 @@ void ParallelConfig::from_flags() {
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_multi_stream_parallel);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(micro_batch_num);
   XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_dp_balance);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(enable_flashcomm1);
+  XLLM_CONFIG_ASSIGN_FROM_FLAG(flashcomm1_sp_min_token_num);
 }
 
 void ParallelConfig::from_json(const JsonReader& json) {
@@ -100,6 +122,8 @@ void ParallelConfig::from_json(const JsonReader& json) {
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_multi_stream_parallel);
   XLLM_CONFIG_ASSIGN_FROM_JSON(micro_batch_num);
   XLLM_CONFIG_ASSIGN_FROM_JSON(enable_dp_balance);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(enable_flashcomm1);
+  XLLM_CONFIG_ASSIGN_FROM_JSON(flashcomm1_sp_min_token_num);
 }
 
 void ParallelConfig::append_config_json(
@@ -124,6 +148,10 @@ void ParallelConfig::append_config_json(
       config_json, default_config, micro_batch_num);
   APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
       config_json, default_config, enable_dp_balance);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, enable_flashcomm1);
+  APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(
+      config_json, default_config, flashcomm1_sp_min_token_num);
 }
 
 ParallelConfig& ParallelConfig::get_instance() {

xllm/core/framework/config/parallel_config.h

@@ -52,7 +52,9 @@ class ParallelConfig final {
          "enable_mm_encoder_dp",
          "enable_multi_stream_parallel",
          "micro_batch_num",
-         "enable_dp_balance"}};
+         "enable_dp_balance",
+         "enable_flashcomm1",
+         "flashcomm1_sp_min_token_num"}};
     return kOptionCategory;
   }
 
@@ -83,6 +85,17 @@ class ParallelConfig final {
 
   PROPERTY(bool, enable_dp_balance) = false;
 
+  // FlashComm1 (sequence parallel) master switch. Defaults to false; when true
+  // and tp>1, the model forward token-shards the residual stream and eligible
+  // RowParallel tails reduce-scatter instead of all-reduce.
+  PROPERTY(bool, enable_flashcomm1) = false;
+
+  // FlashComm1 token gate: SP engages only when dim0 tokens >= this value,
+  // keeping it on for prefill and off for decode where the collective fixed
+  // overhead dominates. 0 disables the gate. Only meaningful with
+  // enable_flashcomm1.
+  PROPERTY(int32_t, flashcomm1_sp_min_token_num) = 1000;
+
   [[nodiscard]] int32_t kv_split_size_effective() const noexcept {
     return kv_split_size_ > 0 ? kv_split_size_ : cp_size_;
   }

xllm/core/framework/parallel_state/CMakeLists.txt

@@ -12,6 +12,7 @@ cc_library(
     rank_generator.h
     parallel_args.h
     parallel_state.h
+    flash_comm1_context.h
     process_group.h
     $<$<BOOL:${USE_NPU}>:npu_rank_table_env.h>
     $<$<BOOL:${USE_NPU}>:npu_process_group.h>
@@ -30,6 +31,7 @@ cc_library(
     dit_mapping.cpp
     parallel_state.cpp
     parallel_state_async.cpp
+    flash_comm1_context.cpp
     process_group.cpp
     $<$<BOOL:${USE_NPU}>:npu_rank_table_env.cpp>
     $<$<BOOL:${USE_NPU}>:npu_process_group.cpp>

xllm/core/framework/parallel_state/flash_comm1_context.cpp

@@ -0,0 +1,52 @@
+/* Copyright 2026 The xLLM Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://github.com/jd-opensource/xllm/blob/main/LICENSE
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "flash_comm1_context.h"
+
+namespace xllm {
+namespace parallel_state {
+namespace {
+
+// Thread-local so a per-thread model forward can publish its own context
+// without cross-thread interference (each worker runs its own forward).
+thread_local FlashComm1Context t_flash_comm1_context;
+
+}  // namespace
+
+const FlashComm1Context& current_flash_comm1_context() {
+  return t_flash_comm1_context;
+}
+
+bool flash_comm1_active() { return t_flash_comm1_context.enabled; }
+
+FlashComm1Guard::FlashComm1Guard(bool enabled,
+                                 int64_t num_tokens,
+                                 int32_t tp_rank,
+                                 int32_t tp_world_size,
+                                 ProcessGroup* tp_group)
+    : previous_(t_flash_comm1_context) {
+  FlashComm1Context context;
+  context.enabled = enabled;
+  context.num_tokens = num_tokens;
+  context.tp_rank = tp_rank;
+  context.tp_world_size = tp_world_size;
+  context.tp_group = tp_group;
+  t_flash_comm1_context = context;
+}
+
+FlashComm1Guard::~FlashComm1Guard() { t_flash_comm1_context = previous_; }
+
+}  // namespace parallel_state
+}  // namespace xllm