From d98057e1661aaf403abe62344f10b602469fd5e4 Mon Sep 17 00:00:00 2001 From: chenchuw886 Date: Thu, 23 Jul 2026 18:42:39 +0800 Subject: [PATCH 1/3] feat(flashcomm1): sequence-parallel infrastructure Add the reusable FlashComm1 (sequence-parallel) primitives, independent of any model: - FlashComm1Context / FlashComm1Guard RAII + current_flash_comm1_context() and flash_comm1_active() in parallel_state, so a decoder stack can publish its per-forward SP state (enabled, num_tokens, tp_rank, tp_world_size, tp_group) to the layers below without threading it through every signature. - Collective helpers in parallel_state: all_gather_dim0_unpad, reduce_scatter_padded_dim0, shard_dim0_padded (dim-0 token sharding with padding for non-divisible token counts). - RowParallelLinear flash_comm1_eligible_ flag: when the active context is enabled and the linear is marked eligible, its all-reduce tail becomes a reduce_scatter_padded_dim0 so the output stays token-sharded. - Unit test for FlashComm1Context/Guard state and activation semantics. --- .../framework/parallel_state/CMakeLists.txt | 14 ++ .../flash_comm1_context_test.cpp | 155 ++++++++++++++++++ .../framework/parallel_state/CMakeLists.txt | 2 + .../parallel_state/flash_comm1_context.cpp | 52 ++++++ .../parallel_state/flash_comm1_context.h | 91 ++++++++++ .../parallel_state/parallel_state.cpp | 94 +++++++++++ .../framework/parallel_state/parallel_state.h | 30 ++++ xllm/core/layers/common/linear.cpp | 17 +- xllm/core/layers/common/linear.h | 14 ++ 9 files changed, 468 insertions(+), 1 deletion(-) create mode 100644 tests/core/framework/parallel_state/flash_comm1_context_test.cpp create mode 100644 xllm/core/framework/parallel_state/flash_comm1_context.cpp create mode 100644 xllm/core/framework/parallel_state/flash_comm1_context.h diff --git a/tests/core/framework/parallel_state/CMakeLists.txt b/tests/core/framework/parallel_state/CMakeLists.txt index 5e13dd3f92..c23ffabcd8 100644 --- a/tests/core/framework/parallel_state/CMakeLists.txt +++ b/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 diff --git a/tests/core/framework/parallel_state/flash_comm1_context_test.cpp b/tests/core/framework/parallel_state/flash_comm1_context_test.cpp new file mode 100644 index 0000000000..35f3f4d615 --- /dev/null +++ b/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 +#include + +#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(), 3.0f); + EXPECT_FLOAT_EQ(shard1[1].item(), 4.0f); + EXPECT_FLOAT_EQ(shard1[2].item(), 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 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(), 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 diff --git a/xllm/core/framework/parallel_state/CMakeLists.txt b/xllm/core/framework/parallel_state/CMakeLists.txt index 26bdab465d..d2705d7ce7 100644 --- a/xllm/core/framework/parallel_state/CMakeLists.txt +++ b/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 $<$:npu_rank_table_env.h> $<$: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 $<$:npu_rank_table_env.cpp> $<$:npu_process_group.cpp> diff --git a/xllm/core/framework/parallel_state/flash_comm1_context.cpp b/xllm/core/framework/parallel_state/flash_comm1_context.cpp new file mode 100644 index 0000000000..b0f7442df7 --- /dev/null +++ b/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 diff --git a/xllm/core/framework/parallel_state/flash_comm1_context.h b/xllm/core/framework/parallel_state/flash_comm1_context.h new file mode 100644 index 0000000000..25a48cff40 --- /dev/null +++ b/xllm/core/framework/parallel_state/flash_comm1_context.h @@ -0,0 +1,91 @@ +/* Copyright 2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#pragma once + +#include + +namespace xllm { + +class ProcessGroup; + +namespace parallel_state { + +// Per-forward FlashComm1 (sequence-parallel) state. +// +// FlashComm1 keeps the residual stream token-sharded across the TP group for +// the whole decoder stack. Each sequence-parallel boundary all-gathers back to +// the full token dimension before attention / MoE, and the RowParallelLinear +// tail all-reduce is replaced by a padded reduce-scatter(dim0) so the output +// returns to a token shard. Norm / hc / gate therefore run on 1/tp_world_size +// tokens. +// +// Layers such as RowParallelLinearImpl::forward() do not receive +// ModelInputParams, so the active context is published via a thread_local set +// by FlashComm1Guard at the top of each model forward(). +struct FlashComm1Context { + // True when a FlashComm1 sequence-parallel forward is active. The caller is + // responsible for folding in the token-count threshold, tp_world_size > 1 and + // any backend / graph checks before setting this. + bool enabled = false; + + // Original (un-sharded, un-padded) number of tokens in this forward. Used to + // unpad after an all-gather at sequence-parallel boundaries. + int64_t num_tokens = 0; + + int32_t tp_rank = 0; + int32_t tp_world_size = 1; + + // TP process group used for the sequence-parallel all-gather / reduce-scatter + // boundaries. Populated by FlashComm1Guard so layers that only see the + // context (e.g. the DSV4 decoder layer, RowParallelLinear) can reach the + // group without extra plumbing. + ProcessGroup* tp_group = nullptr; +}; + +// Returns the FlashComm1 context active on the current thread. When no guard is +// in scope the returned context has enabled == false. +const FlashComm1Context& current_flash_comm1_context(); + +// Convenience helper: true when a FlashComm1 context is active AND enabled. +bool flash_comm1_active(); + +// RAII guard that publishes a FlashComm1 context for the duration of a model +// forward() on the calling thread and restores the previous context on scope +// exit. Guards nest: the previous context is saved and restored, so an inner +// guard (e.g. a nested sub-model forward) does not clobber the outer state. +// +// `enabled` should already fold in the token-count threshold and the +// tp_world_size > 1 / backend checks decided by the caller. +class FlashComm1Guard final { + public: + FlashComm1Guard(bool enabled, + int64_t num_tokens, + int32_t tp_rank, + int32_t tp_world_size, + ProcessGroup* tp_group); + ~FlashComm1Guard(); + + FlashComm1Guard(const FlashComm1Guard&) = delete; + FlashComm1Guard& operator=(const FlashComm1Guard&) = delete; + FlashComm1Guard(FlashComm1Guard&&) = delete; + FlashComm1Guard& operator=(FlashComm1Guard&&) = delete; + + private: + FlashComm1Context previous_; +}; + +} // namespace parallel_state +} // namespace xllm diff --git a/xllm/core/framework/parallel_state/parallel_state.cpp b/xllm/core/framework/parallel_state/parallel_state.cpp index 4f1162859b..2797c21660 100644 --- a/xllm/core/framework/parallel_state/parallel_state.cpp +++ b/xllm/core/framework/parallel_state/parallel_state.cpp @@ -268,6 +268,100 @@ torch::Tensor scatter(torch::Tensor input, return tensor_list[rank]; } +torch::Tensor shard_dim0_padded(const torch::Tensor& input, + int32_t rank, + int32_t world_size) { + if (world_size <= 1) { + return input; + } + CHECK_GE(rank, 0); + CHECK_LT(rank, world_size); + + const int64_t original_dim_size = input.size(0); + const int64_t remainder = original_dim_size % world_size; + const int64_t num_padding = (remainder == 0) ? 0 : (world_size - remainder); + + torch::Tensor padded_input = input; + if (num_padding > 0) { + // Pad only dim0 (rows); leave every other dim untouched. The pad vector is + // ordered from the last dim to the first, two entries (low, high) per dim. + std::vector pad(static_cast(2 * input.dim()), 0); + pad[static_cast(2 * (input.dim() - 1) + 1)] = num_padding; + padded_input = torch::nn::functional::pad( + input, torch::nn::functional::PadFuncOptions(pad)); + } + + const int64_t chunk_size = padded_input.size(0) / world_size; + return padded_input + .slice(0, + static_cast(rank) * chunk_size, + static_cast(rank + 1) * chunk_size) + .contiguous(); +} + +torch::Tensor all_gather_dim0_unpad(const torch::Tensor& input, + ProcessGroup* process_group, + int64_t original_num_tokens) { + if (!process_group) { + return input; + } + const int32_t world_size = process_group->world_size(); + if (world_size == 1) { + return input; + } + // Every rank holds an equal-sized shard, so a plain all-gather along dim0 + // reconstructs the padded full tensor. + const int32_t local_num_tokens = static_cast(input.size(0)); + const std::vector token_num_list( + static_cast(world_size), local_num_tokens); + torch::Tensor full = gather(input, process_group, token_num_list); + if (original_num_tokens >= 0 && full.size(0) > original_num_tokens) { + full = full.slice(0, 0, original_num_tokens).contiguous(); + } + return full; +} + +torch::Tensor reduce_scatter_padded_dim0(const torch::Tensor& input, + ProcessGroup* process_group) { + if (!process_group) { + return input; + } + const int32_t world_size = process_group->world_size(); + if (world_size == 1) { + return input; + } + + const int64_t original_dim_size = input.size(0); + const int64_t remainder = original_dim_size % world_size; + const int64_t num_padding = (remainder == 0) ? 0 : (world_size - remainder); + + torch::Tensor padded_input = input; + if (num_padding > 0) { + // Pad the tail of dim0 (token dim) with zeros for a tensor of arbitrary + // rank. torch pad pairs are ordered from the last dim to the first, so the + // last pair targets dim0. reduce-scatter of a zero-padded tail is a no-op + // for the sum, so correctness is preserved. + std::vector pad(static_cast(2 * input.dim()), 0); + pad[static_cast(2 * (input.dim() - 1) + 1)] = num_padding; + padded_input = torch::nn::functional::pad( + input, torch::nn::functional::PadFuncOptions(pad)); + } + + const int64_t padded_dim_size = padded_input.size(0); + const int64_t chunk_size = padded_dim_size / world_size; + + auto output_shape = padded_input.sizes().vec(); + output_shape[0] = chunk_size; + torch::Tensor output = torch::empty(output_shape, padded_input.options()); + + process_group->reduce_scatter(padded_input, output); + + // NOTE: unlike reduce_scatter(), the trailing padding is intentionally kept so + // that every rank holds an identically-shaped [chunk_size, ...] token shard + // matching shard_dim0_padded. The final all_gather_dim0_unpad strips it. + return output; +} + std::function all_to_all_4D(const torch::Tensor& input, int32_t scatter_idx, int32_t gather_idx, diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h index 905f3c78d9..98084567ce 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -15,6 +15,7 @@ limitations under the License. #pragma once +#include "flash_comm1_context.h" #include "parallel_args.h" #include "process_group.h" @@ -72,6 +73,35 @@ torch::Tensor scatter(torch::Tensor input, ProcessGroup* process_group, int dim = -1); +// FlashComm1 sequence-parallel primitives. +// +// shard_dim0_padded: pad `input` along dim0 up to a multiple of world_size, +// then return this rank's contiguous 1/world_size shard. The padding rows live +// only on the last shard(s); callers restore the original token count with +// all_gather_dim0_unpad. A no-op when process_group is null or world_size == 1. +torch::Tensor shard_dim0_padded(const torch::Tensor& input, + int32_t rank, + int32_t world_size); + +// all_gather_dim0_unpad: all-gather the per-rank shards back into the full +// (padded) token dimension along dim0, then slice off padding so the result +// has exactly `original_num_tokens` rows. When original_num_tokens < 0 the +// gathered (still padded) tensor is returned unsliced. A no-op when +// process_group is null or world_size == 1. +torch::Tensor all_gather_dim0_unpad(const torch::Tensor& input, + ProcessGroup* process_group, + int64_t original_num_tokens); + +// reduce_scatter_padded_dim0: reduce-scatter along dim0, padding dim0 up to a +// multiple of world_size first. Unlike reduce_scatter(), the trailing padding +// is INTENTIONALLY kept so every rank returns an identically-shaped +// [chunk_size, ...] shard that matches shard_dim0_padded (chunk_size = +// ceil(num_tokens / world_size)). This keeps the FlashComm1 sharded residual +// stream uniform across ranks; padding rows are stripped by the final +// all_gather_dim0_unpad. A no-op when process_group is null or world_size == 1. +torch::Tensor reduce_scatter_padded_dim0(const torch::Tensor& input, + ProcessGroup* process_group); + std::function all_to_all_4D(const torch::Tensor& input_, int32_t scatter_idx, int32_t gather_idx, diff --git a/xllm/core/layers/common/linear.cpp b/xllm/core/layers/common/linear.cpp index 8d4612474e..7ce781e951 100644 --- a/xllm/core/layers/common/linear.cpp +++ b/xllm/core/layers/common/linear.cpp @@ -1549,7 +1549,22 @@ torch::Tensor RowParallelLinearImpl::forward(torch::Tensor input) { output = xllm::kernel::matmul(matmul_params); } if (enable_result_reduction_ && world_size_ > 1) { - output = xllm::parallel_state::reduce(output, process_group_); + // FlashComm1: when a sequence-parallel context is active and this layer is + // eligible (e.g. attention o_b_proj), replace the tail all-reduce with a + // padded reduce-scatter(dim0) so the output returns to a token shard that + // matches the sharded residual stream. The collective runs over this + // layer's own TP group, which is the model TP group that produced the + // shard. Otherwise, unchanged all-reduce. + const auto& fc1 = xllm::parallel_state::current_flash_comm1_context(); + if (fc1.enabled && flash_comm1_eligible_) { + // Padded reduce-scatter keeps a uniform [chunk_size, ...] shard on every + // rank (matching shard_dim0_padded), so the sharded residual stream stays + // shape-consistent across ranks. Padding is stripped at the final gather. + output = + xllm::parallel_state::reduce_scatter_padded_dim0(output, process_group_); + } else { + output = xllm::parallel_state::reduce(output, process_group_); + } } return output; } diff --git a/xllm/core/layers/common/linear.h b/xllm/core/layers/common/linear.h index ad3186be2f..a9d73ff90f 100644 --- a/xllm/core/layers/common/linear.h +++ b/xllm/core/layers/common/linear.h @@ -282,6 +282,15 @@ class RowParallelLinearImpl : public torch::nn::Module { } ProcessGroup* process_group() const { return process_group_; } + // Mark this RowParallel layer as eligible for FlashComm1 (sequence parallel). + // When a FlashComm1 context is active, an eligible layer replaces its tail + // all-reduce with a padded reduce-scatter(dim0) so its output returns to a + // token shard. Only layers that feed a token-sharded residual boundary (e.g. + // attention o_b_proj) should be marked. Default false = unchanged behavior. + void set_flash_comm1_eligible(bool eligible) { + flash_comm1_eligible_ = eligible; + } + private: // parameter members, must be registered // we allocate the transpose since linear performs XA^T. @@ -313,6 +322,11 @@ class RowParallelLinearImpl : public torch::nn::Module { // whether to reduce the results bool enable_result_reduction_; + // whether this layer participates in FlashComm1 (sequence parallel). When + // true and a FlashComm1 context is active, the tail all-reduce is replaced by + // a padded reduce-scatter(dim0). + bool flash_comm1_eligible_ = false; + // parallel process group ProcessGroup* process_group_; From 3fd757f4ff259c643d72c56d833611c24a5cd928 Mon Sep 17 00:00:00 2001 From: chenchuw886 Date: Thu, 23 Jul 2026 18:42:50 +0800 Subject: [PATCH 2/3] feat(flashcomm1): wire sequence parallelism into DeepSeek-V4-Flash Drive FlashComm1 from the DeepSeek-V4-Flash decoder stack: - deepseek_v4.h: token-shard the residual stream across the TP group for the whole stack, publish the per-forward FlashComm1Guard, and restore full tokens right before hc_head. Gating policy: engage ONLY in prefill (q_max_seq_len > 1, which also covers chunked-prefill) AND with enough tokens (>= flashcomm1_sp_min_token_num); decode is disabled unconditionally. Empty DP ranks and ACL-graph forwards are excluded. Capture pre_hc_head_hidden_states AFTER the final all-gather so the MTP draft model receives full-token (not sharded) hidden states. - deepseek_v4_decoder_layer.cpp: all-gather the sharded stream to full tokens at the attention and MoE boundaries and re-shard afterwards, so norm / hc / gate run on 1/tp_world_size tokens while attention and MoE see byte-identical full-token input. - deepseek_sparse_attention.cpp: mark o_b_proj as flash_comm1_eligible so its RowParallel tail reduce-scatters back to a token shard. --- .../core/layers/deepseek_v4_decoder_layer.cpp | 39 +++++++++++++ .../npu_torch/deepseek_sparse_attention.cpp | 5 ++ xllm/models/llm/deepseek_v4.h | 58 +++++++++++++++++++ 3 files changed, 102 insertions(+) diff --git a/xllm/core/layers/deepseek_v4_decoder_layer.cpp b/xllm/core/layers/deepseek_v4_decoder_layer.cpp index 8216ae9f07..d86951b3be 100644 --- a/xllm/core/layers/deepseek_v4_decoder_layer.cpp +++ b/xllm/core/layers/deepseek_v4_decoder_layer.cpp @@ -17,6 +17,8 @@ limitations under the License. #include +#include "framework/config/parallel_config.h" +#include "framework/parallel_state/parallel_state.h" #include "kernels/ops_api.h" namespace xllm { @@ -143,11 +145,30 @@ torch::Tensor DeepseekV4DecoderLayerImpl::forward( CHECK(attn_metadata.dsa_metadata) << "DeepseekV4DecoderLayer requires DSA metadata for DSAttention path."; + // FlashComm1 (sequence parallel): x arrives token-sharded across the TP + // group. hc_pre / norm run on the local shard; the full token dimension is + // restored right before attention so the DSA kernels and q/kv projections see + // all tokens. Attention's o_b_proj reduce-scatters the output back to a shard + // (see RowParallelLinearImpl::forward), so residual and attn output stay + // shard-aligned. A no-op when no FlashComm1 context is active. + const auto& fc1 = parallel_state::current_flash_comm1_context(); + + // Gather the token shard back to full tokens (lossless bf16) at an SP + // boundary. + auto fc1_lossless_gather = [&](const torch::Tensor& t) -> torch::Tensor { + return parallel_state::all_gather_dim0_unpad( + t, fc1.tp_group, fc1.num_tokens); + }; + auto residual_attn = x; auto [attn_input, post_attn, comb_attn] = hc_pre(x, hc_attn_fn_, hc_attn_scale_, hc_attn_base_); attn_input = std::get<0>(attn_norm_->forward(attn_input)); + if (fc1.enabled) { + attn_input = fc1_lossless_gather(attn_input); + } + auto& dsa = *(attn_metadata.dsa_metadata); const auto compress_metadata = std::make_tuple( dsa.c1_metadata, dsa.c4_metadata, dsa.c128_metadata, dsa.qli_metadata); @@ -173,6 +194,12 @@ torch::Tensor DeepseekV4DecoderLayerImpl::forward( hc_pre(x, hc_ffn_fn_, hc_ffn_scale_, hc_ffn_base_); ffn_input = std::get<0>(ffn_norm_->forward(ffn_input)); + // MoE boundary: all-gather the shard back to full tokens so gate + MoE run + // with baseline routing semantics; the output is sharded back below. + if (fc1.enabled) { + ffn_input = fc1_lossless_gather(ffn_input); + } + auto ffn_input_2d = ffn_input.reshape({-1, ffn_input.size(-1)}); std::optional gate_input_ids = std::nullopt; if (input_ids.has_value() && input_ids.value().defined()) { @@ -196,6 +223,18 @@ torch::Tensor DeepseekV4DecoderLayerImpl::forward( auto [topk_weights, topk_ids] = gate_->forward(ffn_input_2d, gate_input_ids); ffn_input = moe_mlp_->forward_with_selected_experts( ffn_input, topk_weights, topk_ids, input_params); + + if (fc1.enabled) { + // MoE returns full tokens (expert parallelism shards experts, not tokens); + // shard back so it stays aligned with the sharded residual before hc_post. + const int64_t moe_rows = ffn_input.defined() ? ffn_input.size(0) : 0; + CHECK_EQ(moe_rows, fc1.num_tokens) + << "FlashComm1: MoE output must own this DP replica's full token " + "count before TP token sharding (expected " + << fc1.num_tokens << ", got " << moe_rows << ")"; + ffn_input = parallel_state::shard_dim0_padded( + ffn_input, fc1.tp_rank, fc1.tp_world_size); + } x = hc_post(ffn_input, residual_ffn, post_ffn, comb_ffn); return x; diff --git a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp index 17ebd0e194..4349c9da2e 100644 --- a/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp +++ b/xllm/core/layers/npu_torch/deepseek_sparse_attention.cpp @@ -605,6 +605,11 @@ DSAttentionImpl::DSAttentionImpl(const ModelArgs& args, quant_args, parallel_args.tp_group_, options)); + // FlashComm1: the attention output projection feeds the token-sharded + // residual stream. When a FlashComm1 context is active its tail all-reduce + // becomes a reduce-scatter(dim0), returning the full-token attention output + // to this rank's token shard. + o_b_proj_->set_flash_comm1_eligible(true); } std::tuple> diff --git a/xllm/models/llm/deepseek_v4.h b/xllm/models/llm/deepseek_v4.h index c6e1154e1e..eaed37096a 100644 --- a/xllm/models/llm/deepseek_v4.h +++ b/xllm/models/llm/deepseek_v4.h @@ -34,6 +34,8 @@ limitations under the License. #include "core/framework/config/execution_config.h" #include "core/framework/config/kv_cache_config.h" +#include "core/framework/config/parallel_config.h" +#include "core/framework/parallel_state/parallel_state.h" #include "core/framework/state_dict/utils.h" #include "core/kernels/ops_api.h" #include "core/layers/common/dsa_metadata.h" @@ -431,6 +433,7 @@ class DeepseekV4ModelImpl << ", world_size=" << parallel_args.world_size() << ", dp_size=" << parallel_args.dp_size(); tp_num_heads_ = num_heads_ / dp_local_tp_size_; + tp_group_ = parallel_args.tp_group_; window_size_ = model_args.window_size(); index_n_heads_ = model_args.index_n_heads(); index_head_dim_ = model_args.index_head_dim(); @@ -742,6 +745,50 @@ class DeepseekV4ModelImpl } } + // FlashComm1 (sequence parallel): token-shard the replicated hidden stream + // across the TP group for the whole decoder stack. Each layer all-gathers + // to full tokens before attention / MoE and reduce-scatters back, so + // norm / hc / gate run on 1/tp_world_size tokens. The full token dimension + // is restored right before hc_head. A no-op when --enable_flashcomm1 is + // false. + // + // tp_group_ is the per-DP-replica TP group, so sharding / all-gathering + // stays within one DP replica: the layer all-gathers back to the replica's + // full tokens before the MoE, making MoE input identical to the baseline. + // is_empty_dp_rank agrees across a replica's TP ranks, and replicas own + // disjoint groups, so an empty replica opting out cannot desync others. + // + // Gating: engage only in prefill (q_max_seq_len > 1, which also covers + // chunked-prefill) with num_tokens >= flashcomm1_sp_min_token_num; the + // collectives only amortize with enough tokens. Decode, empty DP ranks and + // ACL-graph forwards are excluded. + const int64_t fc1_num_tokens = h.defined() ? h.size(0) : 0; + const auto& fc1_parallel_cfg = ParallelConfig::get_instance(); + // num_tokens must be >= tp_world_size for a non-degenerate shard; + // sp_min_token_num == 0 keeps only the prefill gate. + const int64_t fc1_sp_min_tokens = std::max( + fc1_parallel_cfg.flashcomm1_sp_min_token_num(), + tp_group_ != nullptr ? tp_group_->world_size() : 1); + const bool fc1_tp_ok = + fc1_parallel_cfg.enable_flashcomm1() && tp_group_ != nullptr && + tp_group_->world_size() > 1 && fc1_num_tokens >= fc1_sp_min_tokens; + const bool fc1_is_prefill = input_params.meta.q_max_seq_len > 1; + const bool fc1_enabled = fc1_tp_ok && fc1_is_prefill && + !acl_graph_forward && !is_empty_dp_rank; + parallel_state::FlashComm1Guard fc1_guard( + fc1_enabled, + fc1_num_tokens, + fc1_enabled ? tp_group_->rank() : 0, + fc1_enabled ? tp_group_->world_size() : 1, + tp_group_); + if (parallel_state::flash_comm1_active()) { + LOG_FIRST_N(INFO, 1) + << "[FlashComm1] active: tp_world_size=" << tp_group_->world_size() + << ", num_tokens=" << fc1_num_tokens; + h = parallel_state::shard_dim0_padded( + h, tp_group_->rank(), tp_group_->world_size()); + } + std::optional residual; for (size_t i = 0; i < layers_.size(); i++) { if (attn_metadata.dsa_metadata) { @@ -820,6 +867,13 @@ class DeepseekV4ModelImpl } #endif } + // FlashComm1: restore the full token dimension (stripping shard padding) + // before hc_head / norm / logits so downstream consumers see all tokens. + // pre_hc_head_hidden_states (exported to MTP as aux_hidden_states) must be + // captured from the full stream, so all-gather first, then capture. + if (parallel_state::flash_comm1_active()) { + h = parallel_state::all_gather_dim0_unpad(h, tp_group_, fc1_num_tokens); + } torch::Tensor pre_hc_head_hidden_states; if (model_args_.num_speculative_tokens() > 0) { pre_hc_head_hidden_states = h; @@ -1546,6 +1600,10 @@ class DeepseekV4ModelImpl int64_t index_topk_ = 512; torch::Device device_{torch::kCPU}; + // FlashComm1 (sequence parallel): the per-DP-replica TP group used for the + // token-shard / all-gather boundaries. Captured at construction. + ProcessGroup* tp_group_ = nullptr; + // DSA cache group info: built once at model init from compress_ratios // caches_info_[layer_id] = vector of DSACacheInfo for each cache in that // layer From 3b3c503223f79bd4357bffe553f263b6f90379f9 Mon Sep 17 00:00:00 2001 From: chenchuw886 Date: Thu, 23 Jul 2026 18:42:56 +0800 Subject: [PATCH 3/3] feat(flashcomm1): config flags and token gate Expose FlashComm1 through configuration: - --enable_flashcomm1 (bool, default false): master switch for DeepSeek-V4-Flash sequence parallelism. - --flashcomm1_sp_min_token_num (int, default 1000): token threshold below which SP stays off even in prefill, since SP only amortizes its shard / all-gather / reduce-scatter fixed overhead with enough tokens. Wired through global_flags.h and ParallelConfig (parallel_config.{h,cpp}). --- xllm/core/common/global_flags.h | 5 ++++ .../core/framework/config/parallel_config.cpp | 28 +++++++++++++++++++ xllm/core/framework/config/parallel_config.h | 15 +++++++++- 3 files changed, 47 insertions(+), 1 deletion(-) diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index a9d65171fd..a9efd46f22 100755 --- a/xllm/core/common/global_flags.h +++ b/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); diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index d29e1e2786..7dc700fc82 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/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() { diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..3330caf881 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/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_; }