diff --git a/CMakeLists.txt b/CMakeLists.txt index 9090501db7..3b6ec5e462 100755 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -500,6 +500,7 @@ if(USE_MLU) $ENV{PYTORCH_MLU_INSTALL_PATH}/../ $ENV{PYTORCH_MLU_INSTALL_PATH}/csrc $ENV{PYTORCH_MLU_INSTALL_PATH}/csrc/include + $ENV{PYTORCH_MLU_OPS_INSTALL_PATH}/include $ENV{NEUWARE_HOME}/include ) @@ -508,6 +509,7 @@ if(USE_MLU) $ENV{PYTORCH_INSTALL_PATH}/lib $ENV{PYTORCH_MLU_INSTALL_PATH}/csrc/lib $ENV{PYTORCH_MLU_INSTALL_PATH} + $ENV{PYTORCH_MLU_OPS_INSTALL_PATH}/lib $ENV{NEUWARE_HOME}/lib64 ) endif() diff --git a/scripts/build_support/env.py b/scripts/build_support/env.py index 452fd87da9..b0f64ac9b7 100644 --- a/scripts/build_support/env.py +++ b/scripts/build_support/env.py @@ -37,6 +37,13 @@ def get_torch_mlu_root_path() -> Optional[str]: except ImportError: return None +def get_torch_mlu_ops_root_path() -> Optional[str]: + try: + import torch_mlu_ops + import os + return os.path.dirname(os.path.abspath(torch_mlu_ops.__file__)) + except ImportError: + return None def get_ixformer_root_path() -> Optional[str]: try: @@ -219,6 +226,7 @@ def set_npu_envs() -> None: def set_mlu_envs() -> None: set_common_envs() os.environ["PYTORCH_MLU_INSTALL_PATH"] = get_torch_mlu_root_path() or "" + os.environ["PYTORCH_MLU_OPS_INSTALL_PATH"] = get_torch_mlu_ops_root_path() or "" def set_cuda_envs() -> None: diff --git a/tests/core/framework/parallel_state/parallel_state_test.cpp b/tests/core/framework/parallel_state/parallel_state_test.cpp index f15aad5712..1a0637464a 100644 --- a/tests/core/framework/parallel_state/parallel_state_test.cpp +++ b/tests/core/framework/parallel_state/parallel_state_test.cpp @@ -86,6 +86,26 @@ struct BroadcastTestParams { int64_t numel; int32_t root_rank; }; + +struct AllGatherAsyncTestParams { + int32_t rank; + int32_t world_size; + int32_t port; + std::string host; + int32_t device_index; + int64_t rows; + int64_t cols; +}; + +struct AllToAllAsyncTestParams { + int32_t rank; + int32_t world_size; + int32_t port; + std::string host; + int32_t device_index; + int64_t chunk; + int64_t cols; +}; // Child process test function int run_reduce_scatter_test_child(const TestParams& params) { try { @@ -350,6 +370,167 @@ int run_allgather_base_test_child(const AllGatherBaseTestParams& params) { return 1; } } + +// Child: launch_all_gather + finish_all_gather over a real group. +// +// Each rank `r` builds an input `arange(rows*cols) + r*1000` (distinct per +// rank so every gather slot is identifiable). `launch_all_gather` must hand +// back a `[world_size, rows, cols]` stack whose slot `s` equals rank `s`'s +// input; `finish_all_gather` returns exactly that stacked tensor. We also check +// the work handle is defined (a real comm was enqueued). +int run_all_gather_async_child(const AllGatherAsyncTestParams& params) { + try { + xllm::Device xllm_device(params.device_index); + xllm_device.set_device(); + torch::Device device = xllm_device.unwrap(); + + auto process_group = create_test_process_group( + params.rank, params.world_size, params.port, params.host, device); + if (!process_group) { + LOG(ERROR) << "all_gather rank " << params.rank << ": PG creation failed"; + return 1; + } + + auto options = torch::TensorOptions() + .dtype(torch::kFloat32) + .device(device) + .requires_grad(false); + auto cpu_options = torch::TensorOptions().dtype(torch::kFloat32); + torch::Tensor input = torch::arange(params.rows * params.cols, options) + .reshape({params.rows, params.cols}) + + static_cast(params.rank * 1000); + + AllGatherAsyncCtx ctx = launch_all_gather(input, process_group.get()); + if (!ctx.work.defined()) { + LOG(ERROR) << "all_gather rank " << params.rank + << ": multi-rank launch did not enqueue a work handle"; + return 1; + } + torch::Tensor stacked = finish_all_gather(ctx); + xllm_device.synchronize_default_stream(); + + if (stacked.dim() != 3 || stacked.size(0) != params.world_size || + stacked.size(1) != params.rows || stacked.size(2) != params.cols) { + LOG(ERROR) << "all_gather rank " << params.rank << ": output shape mismatch"; + return 1; + } + + torch::Tensor stacked_cpu = stacked.to(torch::kCPU); + for (int32_t s = 0; s < params.world_size; ++s) { + torch::Tensor expected = + torch::arange(params.rows * params.cols, cpu_options) + .reshape({params.rows, params.cols}) + + static_cast(s * 1000); + if (!torch::equal(stacked_cpu[s], expected)) { + LOG(ERROR) << "all_gather rank " << params.rank + << ": slot " << s << " mismatch"; + return 1; + } + } + + LOG(INFO) << "all_gather rank " << params.rank << ": passed"; + return 0; + } catch (const std::exception& e) { + LOG(ERROR) << "all_gather rank " << params.rank << ": exception: " << e.what(); + return 1; + } +} + +// Child: launch_all_to_all + finish_all_to_all over a real group. +// +// `launch_all_to_all` does an equal-split all-to-all along dim 0 (input rows +// must be divisible by `world_size`). Standard convention: rank `r` splits its +// `[world_size * chunk, cols]` input into `world_size` blocks of `chunk` rows; +// block `i` is sent to rank `i`. Rank `r`'s output block `i` (rows +// `i*chunk:(i+1)*chunk`) is the block every source rank `s` placed in slot `r` +// of ITS input for destination `r`. +// +// To make the result checkable we encode each input value uniquely as +// `f(dest_rank, src_rank, j)` so every (dest, src, j) triple is distinguishable. +int run_all_to_all_async_child(const AllToAllAsyncTestParams& params) { + try { + xllm::Device xllm_device(params.device_index); + xllm_device.set_device(); + torch::Device device = xllm_device.unwrap(); + + auto process_group = create_test_process_group( + params.rank, params.world_size, params.port, params.host, device); + if (!process_group) { + LOG(ERROR) << "all_to_all rank " << params.rank << ": PG creation failed"; + return 1; + } + + auto cpu_options = torch::TensorOptions().dtype(torch::kFloat32); + + // Build the full per-rank input: block `dest` (rows dest*chunk:(dest+1)*chunk) + // is the payload this rank sends to destination `dest`. Encode each + // value as a unique key: dest*(world_size*cols) + rank*cols + j. + const int64_t total_rows = params.world_size * params.chunk; + torch::Tensor input_cpu = + torch::empty({total_rows, params.cols}, cpu_options); + auto in = input_cpu.accessor(); + for (int32_t dest = 0; dest < params.world_size; ++dest) { + for (int64_t j = 0; j < params.chunk; ++j) { + const int64_t row = static_cast(dest) * params.chunk + j; + for (int64_t c = 0; c < params.cols; ++c) { + in[row][c] = static_cast( + static_cast(dest) * (params.world_size * params.cols) + + static_cast(params.rank) * params.cols + c); + } + } + } + torch::Tensor input = input_cpu.to(device).contiguous(); + + AllToAllAsyncCtx ctx = launch_all_to_all(input, process_group.get()); + if (!ctx.work.defined()) { + LOG(ERROR) << "all_to_all rank " << params.rank + << ": multi-rank launch did not enqueue a work handle"; + return 1; + } + torch::Tensor output = finish_all_to_all(ctx); + xllm_device.synchronize_default_stream(); + + if (output.sizes() != input.sizes()) { + LOG(ERROR) << "all_to_all rank " << params.rank << ": output shape mismatch"; + return 1; + } + + // Expected output on receiver rank: chunk i (rows i*chunk:(i+1)*chunk) of the + // output is the chunk that source rank i sent to destination `rank`: + // output[i*chunk + k][c] = encode(dest=rank, src=i, c) + // = rank*(world_size*cols) + i*cols + c + torch::Tensor output_cpu = output.to(torch::kCPU); + auto out = output_cpu.accessor(); + bool ok = true; + for (int32_t i = 0; i < params.world_size; ++i) { + for (int64_t k = 0; k < params.chunk; ++k) { + const int64_t row = static_cast(i) * params.chunk + k; + for (int64_t c = 0; c < params.cols; ++c) { + float expected = static_cast( + static_cast(params.rank) * + (params.world_size * params.cols) + + static_cast(i) * params.cols + c); + if (std::abs(out[row][c] - expected) > 1e-4f) { + ok = false; + LOG(ERROR) << "all_to_all rank " << params.rank << " row " << row + << " col " << c << ": expected " << expected << " got " + << out[row][c]; + } + } + } + } + if (!ok) { + LOG(ERROR) << "all_to_all rank " << params.rank << ": value mismatch"; + return 1; + } + + LOG(INFO) << "all_to_all rank " << params.rank << ": passed"; + return 0; + } catch (const std::exception& e) { + LOG(ERROR) << "all_to_all rank " << params.rank << ": exception: " << e.what(); + return 1; + } +} // Multi-process test fixture class ReduceScatterMultiDeviceTest : public ::testing::Test { protected: @@ -560,6 +741,149 @@ class BroadcastMultiDeviceTest : public ::testing::Test { int64_t numel_ = 0; }; +// --------------------------------------------------------------------------- +// launch_all_gather / finish_all_gather over a real MLU group +// --------------------------------------------------------------------------- + +class AllGatherAsyncMultiDeviceTest : public ::testing::Test { + protected: + void SetUp() override { + world_size_ = 2; + port_ = 29571; + host_ = "127.0.0.1"; + rows_ = 4; + cols_ = 6; + } + + void RunMultiProcessTest() { + std::vector child_pids; + std::vector child_statuses(world_size_); + + for (int32_t rank = 0; rank < world_size_; ++rank) { + pid_t pid = fork(); + if (pid == 0) { + AllGatherAsyncTestParams params; + params.rank = rank; + params.world_size = world_size_; + params.port = port_; + params.host = host_; + params.device_index = rank % Platform::device_count(); + params.rows = rows_; + params.cols = cols_; + _exit(run_all_gather_async_child(params)); + } else if (pid > 0) { + child_pids.push_back(pid); + } else { + LOG(FATAL) << "Failed to fork child process for rank " << rank; + } + } + + bool all_passed = true; + for (size_t i = 0; i < child_pids.size(); ++i) { + int status; + pid_t waited_pid = waitpid(child_pids[i], &status, 0); + if (waited_pid == child_pids[i]) { + if (WIFEXITED(status)) { + child_statuses[i] = WEXITSTATUS(status); + if (child_statuses[i] != 0) { + all_passed = false; + LOG(ERROR) << "Child process for rank " << i << " exited with code " + << child_statuses[i]; + } + } else { + all_passed = false; + LOG(ERROR) << "Child process for rank " << i + << " did not exit normally"; + } + } else { + all_passed = false; + LOG(ERROR) << "Failed to wait for child process for rank " << i; + } + } + + CHECK(all_passed) << "One or more child processes failed"; + } + + int32_t world_size_ = 0; + int32_t port_ = 0; + std::string host_; + int64_t rows_ = 0; + int64_t cols_ = 0; +}; + +// --------------------------------------------------------------------------- +// launch_all_to_all / finish_all_to_all over a real MLU group +// --------------------------------------------------------------------------- + +class AllToAllAsyncMultiDeviceTest : public ::testing::Test { + protected: + void SetUp() override { + world_size_ = 2; + port_ = 29581; + host_ = "127.0.0.1"; + // Equal-split block size (rows sent to / received from each rank). Picked + // == world_size so the per-rank input has world_size*chunk rows and the + // transpose is a clean [world_size, chunk, cols] dim-(0,1) swap. + chunk_ = world_size_; + cols_ = 3; + } + + void RunMultiProcessTest() { + std::vector child_pids; + std::vector child_statuses(world_size_); + + for (int32_t rank = 0; rank < world_size_; ++rank) { + pid_t pid = fork(); + if (pid == 0) { + AllToAllAsyncTestParams params; + params.rank = rank; + params.world_size = world_size_; + params.port = port_; + params.host = host_; + params.device_index = rank % Platform::device_count(); + params.chunk = chunk_; + params.cols = cols_; + _exit(run_all_to_all_async_child(params)); + } else if (pid > 0) { + child_pids.push_back(pid); + } else { + LOG(FATAL) << "Failed to fork child process for rank " << rank; + } + } + + bool all_passed = true; + for (size_t i = 0; i < child_pids.size(); ++i) { + int status; + pid_t waited_pid = waitpid(child_pids[i], &status, 0); + if (waited_pid == child_pids[i]) { + if (WIFEXITED(status)) { + child_statuses[i] = WEXITSTATUS(status); + if (child_statuses[i] != 0) { + all_passed = false; + LOG(ERROR) << "Child process for rank " << i << " exited with code " + << child_statuses[i]; + } + } else { + all_passed = false; + LOG(ERROR) << "Child process for rank " << i + << " did not exit normally"; + } + } else { + all_passed = false; + LOG(ERROR) << "Failed to wait for child process for rank " << i; + } + } + + CHECK(all_passed) << "One or more child processes failed"; + } + + int32_t world_size_ = 0; + int32_t port_ = 0; + std::string host_; + int64_t chunk_ = 0; + int64_t cols_ = 0; +}; + TEST_F(BroadcastMultiDeviceTest, BroadcastFromRoot0) { RunMultiProcessTest(/*root_rank=*/0); } @@ -586,6 +910,14 @@ TEST_F(ReduceScatterMultiDeviceTest, LargeInputTest) { TEST_F(AllGatherBaseMultiDeviceTest, BasicTest) { RunMultiProcessTest(input_size_, hidden_dim_); } + +TEST_F(AllGatherAsyncMultiDeviceTest, StacksPerRankInputs) { + RunMultiProcessTest(); +} + +TEST_F(AllToAllAsyncMultiDeviceTest, EqualSplitBlockSwap) { + RunMultiProcessTest(); +} } // namespace test } // namespace parallel_state } // namespace xllm diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 45b2f9f1aa..ab922eac8f 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -105,6 +105,8 @@ DECLARE_int32(ep_size); DECLARE_int32(cp_size); +DECLARE_int32(dcp_size); + DECLARE_int64(tp_size); DECLARE_int64(sp_size); diff --git a/xllm/core/common/options.cpp b/xllm/core/common/options.cpp index bfe237b1f3..4e796f9fe4 100644 --- a/xllm/core/common/options.cpp +++ b/xllm/core/common/options.cpp @@ -56,6 +56,7 @@ std::string Options::to_string() const { << ", enable_chunked_prefill: " << enable_chunked_prefill() << ", cp_size: " << cp_size() << ", master_node_addr: " << master_node_addr().value_or("null") + << ", dcp_size: " << dcp_size() << ", instance_role: " << instance_role().to_string() << ", transfer_listen_port: " << transfer_listen_port() << ", nnodes: " << nnodes() << ", node_rank: " << node_rank() diff --git a/xllm/core/common/options.h b/xllm/core/common/options.h index e1c56661fc..bfaa2a52fe 100644 --- a/xllm/core/common/options.h +++ b/xllm/core/common/options.h @@ -128,6 +128,10 @@ class Options { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, dcp_size) = 1; + + PROPERTY(int32_t, cp_kv_cache_interleave_size) = 0; + PROPERTY(int32_t, ep_size) = 1; PROPERTY(int32_t, tp_size) = 1; diff --git a/xllm/core/distributed_runtime/master.cpp b/xllm/core/distributed_runtime/master.cpp index b1743ddde2..513426e4c0 100644 --- a/xllm/core/distributed_runtime/master.cpp +++ b/xllm/core/distributed_runtime/master.cpp @@ -74,8 +74,12 @@ std::optional validate_model_cp(const Options& options, if (options.cp_size() < 1) { return "cp_size must be greater than or equal to 1"; } + if (options.dcp_size() < 1) { + return "dcp_size must be greater than or equal to 1"; + } const bool use_model_partition = options.cp_size() > 1 && Platform::uses_model_cp_partition(); + const bool use_dcp = options.dcp_size() > 1 && Platform::uses_decode_dcp(); if (!use_model_partition) { return std::nullopt; } @@ -95,8 +99,9 @@ std::optional validate_model_cp(const Options& options, return "MLU CP does not support model_type=" + model_type; } if (options.instance_role() != InstanceRole::DEFAULT && - options.instance_role() != InstanceRole::PREFILL) { - return "MLU CP supports only DEFAULT or PREFILL roles"; + options.instance_role() != InstanceRole::PREFILL && + options.instance_role() != InstanceRole::DECODE) { + return "MLU CP supports only DEFAULT, PREFILL, or DECODE roles"; } if (options.dp_size() != 1) { return "MLU CP requires dp_size == 1"; @@ -213,7 +218,10 @@ Master::Master(const Options& options, EngineType type) // World size is the node count (one worker per process). const int32_t global_world_size = options_.nnodes(); std::string cp_model_type; - if (options_.cp_size() > 1 && Platform::uses_model_cp_partition()) { + const bool needs_cp_model_type = + (options_.cp_size() > 1 && Platform::uses_model_cp_partition()) || + (options_.dcp_size() > 1 && Platform::uses_decode_dcp()); + if (needs_cp_model_type) { cp_model_type = util::get_model_type(model_path, options_.backend()); } const std::optional cp_error = @@ -234,11 +242,20 @@ Master::Master(const Options& options, EngineType type) print_startup_banner(model_path, options_.backend(), options_.node_rank()); LOG(INFO) << "Master init options: " << options_.to_string(); ParallelConfig::get_instance().cp_size(options_.cp_size()); + ParallelConfig::get_instance().dcp_size(options_.dcp_size()); + ParallelConfig::get_instance().cp_kv_cache_interleave_size( + options_.cp_kv_cache_interleave_size()); + const bool use_dcp = options_.dcp_size() > 1 && Platform::uses_decode_dcp(); + if (use_dcp) { + ParallelConfig::get_instance().kv_split_size(options_.dcp_size()); + } const char* cp_partition_stage = options_.cp_size() <= 1 ? "disabled" : (Platform::uses_model_cp_partition() ? "model" : "worker"); LOG(INFO) << "Resolved CP config: cp_size=" << options_.cp_size() + << ", dcp_size=" << options_.dcp_size() << ", kv_split_size=" + << ParallelConfig::get_instance().kv_split_size_effective() << ", world_size=" << global_world_size << ", dp_size=" << options_.dp_size() << ", ep_size=" << options_.ep_size() @@ -315,6 +332,8 @@ Master::Master(const Options& options, EngineType type) .task_type(options.task_type()) .enable_mla(options_.enable_mla()) .cp_size(options_.cp_size()) + .dcp_size(options_.dcp_size()) + .cp_kv_cache_interleave_size(options_.cp_kv_cache_interleave_size()) .npu_kernel_backend(options_.npu_kernel_backend()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .enable_offline_inference(options_.enable_offline_inference()) @@ -389,6 +408,8 @@ Master::Master(const Options& options, EngineType type) .dp_size(options.dp_size()) .ep_size(options.ep_size()) .cp_size(options_.cp_size()) + .dcp_size(options_.dcp_size()) + .cp_kv_cache_interleave_size(options_.cp_kv_cache_interleave_size()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .max_tokens_per_batch(options_.max_tokens_per_batch()) .max_seqs_per_batch(options_.max_seqs_per_batch()) @@ -443,6 +464,8 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) + .dcp_size(options_.dcp_size()) + .cp_kv_cache_interleave_size(options_.cp_kv_cache_interleave_size()) .enable_chunked_prefill(options_.enable_chunked_prefill()) .max_tokens_per_batch(options_.max_tokens_per_batch()) .max_seqs_per_batch(options_.max_seqs_per_batch()) @@ -509,6 +532,8 @@ Master::Master(const Options& options, EngineType type) .dp_size(options_.dp_size()) .ep_size(options_.ep_size()) .cp_size(options_.cp_size()) + .dcp_size(options_.dcp_size()) + .cp_kv_cache_interleave_size(options_.cp_kv_cache_interleave_size()) .max_seqs_per_batch(options_.max_seqs_per_batch()) .beam_width(options_.beam_width()) .max_tokens_per_batch(options_.max_tokens_per_batch()) diff --git a/xllm/core/framework/config/parallel_config.cpp b/xllm/core/framework/config/parallel_config.cpp index d29e1e2786..3c5fbaf356 100644 --- a/xllm/core/framework/config/parallel_config.cpp +++ b/xllm/core/framework/config/parallel_config.cpp @@ -26,6 +26,21 @@ DEFINE_int32(ep_size, 1, "Expert parallel size for MoE model."); DEFINE_int32(cp_size, 1, "Context parallel size for DSA attention."); +DEFINE_int32(dcp_size, + 1, + "Decode context parallel size. DCP shards decode KV cache across " + "dcp_size ranks by reusing the kv_split_size mechanism (the engine " + "auto-sets kv_split_size = dcp_size so each rank owns a 1/dcp " + "row-band within every logical block, an interleaved subsampling " + "of the tokens). Phase 1 requires tp_size % dcp_size == 0."); + +DEFINE_int32(cp_kv_cache_interleave_size, + 0, + "Interleave size of KV cache storage while using CP. " + "0=auto (uses block_size, block-banded ownership, backward " + "compatible); 1=token-level interleaving; N>0=interleave at N " + "tokens. Must divide block_size."); + DEFINE_int32(kv_split_size, 1, "KV-cache split width. 0 falls back to cp_size (legacy); 1 means " @@ -75,6 +90,8 @@ void ParallelConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(ep_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(cp_size); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dcp_size); + XLLM_CONFIG_ASSIGN_FROM_FLAG(cp_kv_cache_interleave_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(kv_split_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(tp_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(sp_size); @@ -91,6 +108,8 @@ void ParallelConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(ep_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cp_size); + XLLM_CONFIG_ASSIGN_FROM_JSON(dcp_size); + XLLM_CONFIG_ASSIGN_FROM_JSON(cp_kv_cache_interleave_size); XLLM_CONFIG_ASSIGN_FROM_JSON(tp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(sp_size); XLLM_CONFIG_ASSIGN_FROM_JSON(cfg_size); @@ -108,6 +127,10 @@ void ParallelConfig::append_config_json( APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, dp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, ep_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, cp_size); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dcp_size); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, cp_kv_cache_interleave_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, tp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT(config_json, default_config, sp_size); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/parallel_config.h b/xllm/core/framework/config/parallel_config.h index 22159a1c2a..73f6c83f97 100644 --- a/xllm/core/framework/config/parallel_config.h +++ b/xllm/core/framework/config/parallel_config.h @@ -44,6 +44,8 @@ class ParallelConfig final { {"dp_size", "ep_size", "cp_size", + "dcp_size", + "cp_kv_cache_interleave_size", "tp_size", "sp_size", "cfg_size", @@ -62,6 +64,10 @@ class ParallelConfig final { PROPERTY(int32_t, cp_size) = 1; + PROPERTY(int32_t, dcp_size) = 1; + + PROPERTY(int32_t, cp_kv_cache_interleave_size) = 0; + // 0 means follow cp_size (legacy KV-split width). PROPERTY(int32_t, kv_split_size) = 1; @@ -86,6 +92,12 @@ class ParallelConfig final { [[nodiscard]] int32_t kv_split_size_effective() const noexcept { return kv_split_size_ > 0 ? kv_split_size_ : cp_size_; } + + [[nodiscard]] int32_t cp_kv_cache_interleave_size_effective( + int32_t block_size) const noexcept { + return cp_kv_cache_interleave_size_ > 0 ? cp_kv_cache_interleave_size_ + : block_size; + } }; } // namespace xllm diff --git a/xllm/core/framework/parallel_state/collective_communicator.cpp b/xllm/core/framework/parallel_state/collective_communicator.cpp index ab8eb11e2e..8686da21a2 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.cpp +++ b/xllm/core/framework/parallel_state/collective_communicator.cpp @@ -226,6 +226,10 @@ CollectiveCommunicator::CollectiveCommunicator(int global_rank, global_rank, world_size, dp_size, cp_size, nullptr, ep_size); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().dcp_size()); + parallel_args_->cp_kv_cache_interleave_size( + ::xllm::ParallelConfig::get_instance().cp_kv_cache_interleave_size()); return; } @@ -283,11 +287,19 @@ CollectiveCommunicator::CollectiveCommunicator(int global_rank, dispatchAndCombineHcclComm); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().dcp_size()); + parallel_args_->cp_kv_cache_interleave_size( + ::xllm::ParallelConfig::get_instance().cp_kv_cache_interleave_size()); #else parallel_args_ = std::make_unique( global_rank, world_size, dp_size, cp_size, nullptr, ep_size); parallel_args_->kv_split_size( ::xllm::ParallelConfig::get_instance().kv_split_size()); + parallel_args_->dcp_size( + ::xllm::ParallelConfig::get_instance().dcp_size()); + parallel_args_->cp_kv_cache_interleave_size( + ::xllm::ParallelConfig::get_instance().cp_kv_cache_interleave_size()); #endif } @@ -391,6 +403,37 @@ void CollectiveCommunicator::create_process_groups( parallel_args_->cp_group_ = tp_group_.get(); port += dp_size + single_rank_group_port_gap + single_rank_group_count; + const int32_t dcp_size = parallel_args_->dcp_size_effective(); + if (dcp_size > 1) { + CHECK_EQ(tp_size % dcp_size, 0) + << "DCP requires tp_size % dcp_size == 0 (tp_size=" << tp_size + << ", dcp_size=" << dcp_size << ")"; + const int32_t dcp_group_size = dcp_size; + const int32_t tp_rank = global_rank % tp_size; + const int32_t dcp_rank = tp_rank % dcp_size; + const int32_t dcp_group_base = + (tp_rank / dcp_size) * dcp_size; + std::vector dcp_group_ranks; + dcp_group_ranks.reserve(dcp_group_size); + for (int32_t member = 0; member < dcp_group_size; ++member) { + const int32_t member_tp_rank = dcp_group_base + member; + const int32_t dp_base = global_rank - tp_rank; + dcp_group_ranks.push_back(dp_base + member_tp_rank); + } + port += 1; + dcp_group_ = create_process_group(global_rank, + dcp_rank, + dcp_group_ranks, + world_size, + dcp_group_size, + port + global_rank / tp_size + 1, + host, + "dcp_group", + device); + parallel_args_->dcp_group_ = dcp_group_.get(); + port += tp_size; + } + if (dp_size > 1) { port_offset = global_rank % tp_size + 1; dp_local_process_group_ = create_process_group(global_rank, diff --git a/xllm/core/framework/parallel_state/collective_communicator.h b/xllm/core/framework/parallel_state/collective_communicator.h index d29dd31f43..b3170887d9 100644 --- a/xllm/core/framework/parallel_state/collective_communicator.h +++ b/xllm/core/framework/parallel_state/collective_communicator.h @@ -45,6 +45,7 @@ class CollectiveCommunicator : public CollectiveCommunicatorBase { // model-side CP path aliases ParallelArgs::cp_group_ to tp_group_ instead of // constructing a separate communicator, so this stays empty for now. std::unique_ptr cp_group_; + std::unique_ptr dcp_group_; std::unique_ptr moe_tp_group_; std::unique_ptr moe_ep_group_; }; diff --git a/xllm/core/framework/parallel_state/parallel_args.h b/xllm/core/framework/parallel_state/parallel_args.h index 8203680cd6..60ff87c58a 100644 --- a/xllm/core/framework/parallel_state/parallel_args.h +++ b/xllm/core/framework/parallel_state/parallel_args.h @@ -148,6 +148,11 @@ struct ParallelArgs { // cp size PROPERTY(int32_t, cp_size) = 1; + // dcp size + PROPERTY(int32_t, dcp_size) = 1; + + PROPERTY(int32_t, cp_kv_cache_interleave_size) = 0; + // Derived: CP rank of the current process within its DP group. // rank layout: dp_rank * (cp_size * tp_size) + cp_rank * tp_size + tp_rank [[nodiscard]] int32_t cp_rank() const noexcept { @@ -158,6 +163,28 @@ struct ParallelArgs { return (rank_ % (cp_size_ * tp_sz)) / tp_sz; } + [[nodiscard]] int32_t dcp_size_effective() const noexcept { + return dcp_size_ > 0 ? dcp_size_ : 1; + } + + [[nodiscard]] int32_t cp_kv_cache_interleave_size_effective( + int32_t block_size) const noexcept { + return cp_kv_cache_interleave_size_ > 0 ? cp_kv_cache_interleave_size_ + : block_size; + } + + [[nodiscard]] int32_t dcp_rank() const noexcept { + if (dcp_size_ <= 1) { + return 0; + } + int32_t tp_sz = world_size_ / dp_size_ / cp_size_; + if (tp_sz <= 0 || tp_sz < dcp_size_) { + return 0; + } + int32_t tp_rank = (rank_ % tp_sz); + return tp_rank / (tp_sz / dcp_size_); + } + // KV-cache split width. 0 == "follow cp_size" (legacy). Use // `kv_split_size_effective()` instead of reading the raw value when computing // strides / block sizes; the raw setter is kept so the engine can override @@ -223,6 +250,7 @@ struct ParallelArgs { // rank set, so this temporarily aliases the TP group. Keep a distinct handle // for a future orthogonal CP x TP topology with a standalone CP group. ProcessGroup* cp_group_ = nullptr; + ProcessGroup* dcp_group_ = nullptr; ProcessGroup* moe_ep_group_ = nullptr; ProcessGroup* moe_tp_group_ = nullptr; diff --git a/xllm/core/framework/parallel_state/parallel_state.h b/xllm/core/framework/parallel_state/parallel_state.h index 905f3c78d9..788de0c51c 100644 --- a/xllm/core/framework/parallel_state/parallel_state.h +++ b/xllm/core/framework/parallel_state/parallel_state.h @@ -39,6 +39,18 @@ struct ReduceAsyncCtx { c10::intrusive_ptr work; }; +struct AllGatherAsyncCtx { + torch::Tensor input; + torch::Tensor stacked; + c10::intrusive_ptr work; +}; + +struct AllToAllAsyncCtx { + torch::Tensor input; + torch::Tensor output; + c10::intrusive_ptr work; +}; + std::optional get_dp_attn_parallel_args( const ParallelArgs& parallel_args); @@ -60,6 +72,16 @@ ReduceAsyncCtx launch_reduce(torch::Tensor input, ProcessGroup* process_group); torch::Tensor finish_reduce(ReduceAsyncCtx ctx); +AllGatherAsyncCtx launch_all_gather(const torch::Tensor& input, + ProcessGroup* process_group); + +torch::Tensor finish_all_gather(AllGatherAsyncCtx ctx); + +AllToAllAsyncCtx launch_all_to_all(const torch::Tensor& input, + ProcessGroup* process_group); + +torch::Tensor finish_all_to_all(AllToAllAsyncCtx ctx); + torch::Tensor all_gather_interleaved(const torch::Tensor& input, ProcessGroup* process_group); diff --git a/xllm/core/framework/parallel_state/parallel_state_async.cpp b/xllm/core/framework/parallel_state/parallel_state_async.cpp index 1d136ca5d4..0ca91a0886 100644 --- a/xllm/core/framework/parallel_state/parallel_state_async.cpp +++ b/xllm/core/framework/parallel_state/parallel_state_async.cpp @@ -81,5 +81,59 @@ ReduceAsyncCtx launch_reduce(torch::Tensor input, ProcessGroup* process_group) { return ctx; } +AllGatherAsyncCtx launch_all_gather(const torch::Tensor& input, + ProcessGroup* process_group) { + AllGatherAsyncCtx ctx; + if (!process_group || process_group->world_size() <= 1) { + ctx.input = input.contiguous(); + ctx.stacked = ctx.input.unsqueeze(0); + return ctx; + } + + ctx.input = input.contiguous(); + auto stacked_shape = ctx.input.sizes().vec(); + stacked_shape.insert(stacked_shape.begin(), process_group->world_size()); + ctx.stacked = torch::empty(stacked_shape, ctx.input.options()); + ctx.work = process_group->allgather_base_async(ctx.input, ctx.stacked); + return ctx; +} + +torch::Tensor finish_all_gather(AllGatherAsyncCtx ctx) { + if (ctx.work.defined()) { + ctx.work->wait(); + } + return ctx.stacked; +} + +AllToAllAsyncCtx launch_all_to_all(const torch::Tensor& input, + ProcessGroup* process_group) { + AllToAllAsyncCtx ctx; + if (!process_group || process_group->world_size() <= 1) { + ctx.input = input.contiguous(); + ctx.output = ctx.input; + return ctx; + } + + const int32_t world_size = process_group->world_size(); + ctx.input = input.contiguous(); + CHECK_EQ(ctx.input.size(0) % world_size, 0) + << "launch_all_to_all: input dim 0 (" << ctx.input.size(0) + << ") must be divisible by world_size (" << world_size << ")"; + ctx.output = torch::empty_like(ctx.input); + process_group->all_to_all_single(ctx.output, + ctx.input, + /*output_split_sizes=*/{}, + /*input_split_sizes=*/{}, + /*async_op=*/true, + &ctx.work); + return ctx; +} + +torch::Tensor finish_all_to_all(AllToAllAsyncCtx ctx) { + if (ctx.work.defined()) { + ctx.work->wait(); + } + return ctx.output; +} } // namespace parallel_state } // namespace xllm diff --git a/xllm/core/layers/mlu/CMakeLists.txt b/xllm/core/layers/mlu/CMakeLists.txt index 84ea08bfc8..b1347e19a6 100755 --- a/xllm/core/layers/mlu/CMakeLists.txt +++ b/xllm/core/layers/mlu/CMakeLists.txt @@ -33,6 +33,7 @@ cc_library( SRCS attention.cpp deepseek_v2_attention.cpp + deepseek_v2_attention_dcp.cpp deepseek_v2_attention_sp.cpp deepseek_v2_decoder_layer_impl.cpp deepseek_v2_sparse_moe_block.cpp diff --git a/xllm/core/layers/mlu/deepseek_v2_attention.cpp b/xllm/core/layers/mlu/deepseek_v2_attention.cpp index 46889d9a08..84884fd425 100644 --- a/xllm/core/layers/mlu/deepseek_v2_attention.cpp +++ b/xllm/core/layers/mlu/deepseek_v2_attention.cpp @@ -43,6 +43,11 @@ DeepseekV2AttentionImpl::DeepseekV2AttentionImpl( has_indexer_ = enable_lighting_indexer_ && enable_indexer; use_full_replicated_attention_weights_ = parallel_args.cp_size() > 1 && Platform::uses_model_cp_partition(); + dcp_enabled_ = parallel_args.dcp_size() > 1 && + parallel_args.dcp_group_ != nullptr; + dcp_group_ = parallel_args.dcp_group_; + dcp_size_ = parallel_args.dcp_size_effective(); + kv_split_rank_ = parallel_args.kv_split_rank(); const int64_t tp_size = parallel_args.tp_group_->world_size(); int64_t hidden_size = args.hidden_size(); int64_t num_heads = args.n_heads(); @@ -54,6 +59,8 @@ DeepseekV2AttentionImpl::DeepseekV2AttentionImpl( tp_heads_ = {num_heads / tp_size, num_heads / tp_size}; full_heads_ = {num_heads, num_heads}; float scaling = std::pow(qk_head_dim_, -0.5f); + sliding_window_ = args.sliding_window(); + attn_scale_ = scaling; ProcessGroup* weight_group = use_replicated_attn_weights() && @@ -438,6 +445,11 @@ torch::Tensor DeepseekV2AttentionImpl::forward( DsaTopkTransfer* topk_transfer) { bool is_prefill_or_chunked_prefill = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + if (dcp_enabled_ && !is_prefill_or_chunked_prefill) { + return forward_dcp(positions, hidden_states, attn_metadata, kv_cache, + topk_transfer); + } + if (sp_ctx != nullptr && can_use_sp()) { // DSA cross-layer top-k sharing under sequence parallel is out of scope // for the eager path; the sequence parallel branch always recomputes. diff --git a/xllm/core/layers/mlu/deepseek_v2_attention.h b/xllm/core/layers/mlu/deepseek_v2_attention.h index 9fce9a6e7c..e2d9e12ebc 100644 --- a/xllm/core/layers/mlu/deepseek_v2_attention.h +++ b/xllm/core/layers/mlu/deepseek_v2_attention.h @@ -178,6 +178,12 @@ class DeepseekV2AttentionImpl : public torch::nn::Module { return use_replicated_attn_weights() ? full_heads_ : tp_heads_; } + torch::Tensor forward_dcp(const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + DsaTopkTransfer* topk_transfer = nullptr); + private: bool use_full_replicated_attention_weights_ = false; bool use_fused_mla_qkv_ = false; @@ -185,6 +191,10 @@ class DeepseekV2AttentionImpl : public torch::nn::Module { bool has_indexer_ = false; bool has_trans_ = false; bool interleaved_ = false; + bool dcp_enabled_ = false; + int64_t sliding_window_ = -1; + int32_t kv_split_rank_ = 0; + float attn_scale_ = 1.0f; double eps_; int64_t qk_head_dim_; int64_t v_head_dim_; @@ -215,6 +225,8 @@ class DeepseekV2AttentionImpl : public torch::nn::Module { std::shared_ptr indexer_rotary_emb_; Indexer indexer_{nullptr}; std::unique_ptr sp_comm_stream_; + ProcessGroup* dcp_group_ = nullptr; + int32_t dcp_size_ = 1; }; TORCH_MODULE(DeepseekV2Attention); diff --git a/xllm/core/layers/mlu/deepseek_v2_attention_dcp.cpp b/xllm/core/layers/mlu/deepseek_v2_attention_dcp.cpp new file mode 100644 index 0000000000..e3ae8662ee --- /dev/null +++ b/xllm/core/layers/mlu/deepseek_v2_attention_dcp.cpp @@ -0,0 +1,234 @@ +/* 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 "deepseek_v2_attention.h" + +#include +#include + +#include "framework/parallel_state/parallel_state.h" +#include "kernels/mlu/mlu_ops_api.h" +#include "kernels/ops_api.h" +#include "platform/platform.h" + +namespace xllm { +namespace layer { + +namespace { +void check_phase1_dcp_geometry(int32_t dcp_size, + int64_t tp_heads_attn, + int64_t full_heads_attn) { + CHECK_EQ(dcp_size * tp_heads_attn, full_heads_attn) + << "Phase 1 DCP requires dcp_size * tp_heads == full_heads (i.e. " + "tp_size == dcp_size); tp_size > dcp_size is a Phase 2 task."; +} + +} // namespace + +torch::Tensor DeepseekV2AttentionImpl::forward_dcp( + const torch::Tensor& positions, + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + DsaTopkTransfer* topk_transfer) { + CHECK_GT(dcp_size_, 1) << "forward_dcp requires dcp_size_ > 1."; + + check_phase1_dcp_geometry(dcp_size_, tp_heads_.attn, full_heads_.attn); + + const int64_t tokens = hidden_states.size(0); + auto k_cache = kv_cache.get_k_cache(); + auto k_cache_scale = kv_cache.get_k_cache_scale(); + auto query_prep = prep_query(hidden_states, tp_heads_); + torch::Tensor q_input = torch::empty( + {tokens, tp_heads_.attn, kv_lora_rank_ + qk_rope_head_dim_}, + hidden_states.options()); + torch::Tensor latent_cache = kv_a_proj_with_mqa_(hidden_states); + fill_q_input(q_input, + query_prep.q, + positions, + attn_metadata, + /*use_prompt_rope=*/false); + decode_kv_pre_base(latent_cache, positions, attn_metadata, /*use_prompt_rope=*/false); + + // DSA cross-layer top-k sharing. A Shared layer reuses the previous Full + // layer's sparse block table (skipping indexer recompute); an Output layer + // exports its freshly computed block table so the next Shared layer can + // reuse it. Under DCP the indexer (with the slot recovery below) produces a + // rank-consistent global top-k, so the published state is safe to share + // across ranks just like the eager path. + std::optional shared_block_tables; + std::optional shared_context_lens; + const bool has_shared_topk = + topk_transfer != nullptr && topk_transfer->input() != nullptr; + if (has_shared_topk) { + shared_block_tables = topk_transfer->input()->block_tables(); + shared_context_lens = topk_transfer->input()->context_lens(); + } + + // DCP indexer index-cache slot recovery. + // The worker masks `new_cache_slots` (== attn_metadata.slot_mapping) by + // `pos % dcp_size == dcp_rank` so the sharded MAIN MLA KV cache only receives + // this rank's 1/dcp tokens. The INDEX cache, however, is a full replica on + // every rank (KVCacheShape::init_index_cache_shape does not shard by + // world_size, and the indexer weights are replicated), so it must receive ALL + // new tokens -- otherwise the indexer's top-k selection runs against a + // 1/dcp-populated cache and produces wrong block tables. Each rank keeps the + // original slot id for its owned tokens and -1 elsewhere, so an AllGather + + // max over the dcp group recovers the full (unmasked) slot map. The indexer + // then writes the full-replica index cache with this map, while the main KV + // write below keeps the masked (sharded) map. A Shared layer skips the + // indexer (it reuses the prior top-k), so the recovery AllGather is skipped + // too -- only Full/plain layers that actually recompute the indexer pay for + // it. + torch::Tensor masked_slot_mapping = attn_metadata.slot_mapping; + AttentionMetadata index_attn_metadata = attn_metadata; + const bool indexer_will_run = + enable_lighting_indexer_ && !has_shared_topk; + if (indexer_will_run && masked_slot_mapping.defined() && + dcp_group_ != nullptr) { + auto slot_gather_ctx = parallel_state::launch_all_gather( + masked_slot_mapping, dcp_group_); + torch::Tensor stacked_slots = + parallel_state::finish_all_gather(std::move(slot_gather_ctx)); + stacked_slots = stacked_slots.view({dcp_size_, -1}); + index_attn_metadata.slot_mapping = + std::get<0>(torch::max(stacked_slots, /*dim=*/0, /*keepdim=*/false)); + } + + AttentionMetadata local_meta = build_mla_attention_metadata( + positions, + hidden_states, + query_prep.q_norm, + latent_cache, + index_attn_metadata, + kv_cache, + k_cache_scale, + /*is_prefill_phase=*/false, + /*slot_mapping=*/std::nullopt, + shared_block_tables, + shared_context_lens); + + // An Output layer exports the freshly computed sparse block table (held in + // the resolved metadata) so the caller can cache it for the next Shared + // layer. finish_layer CHECKs that a producer publishes, so this must run + // before the main KV write below. + if (topk_transfer != nullptr && topk_transfer->captures_output()) { + topk_transfer->publish_output(DsaTopkState( + local_meta.block_table, local_meta.kv_seq_lens)); + } + + { + torch::Tensor key = latent_cache.unsqueeze(1); + xllm::kernel::ReshapePagedCacheParams params; + params.key = key; + params.k_cache = k_cache; + // Main MLA KV is DCP-sharded: write only this rank's owned tokens. Slots + // == -1 (non-owned) are skipped by reshape_paged_cache / + // quant_to_paged_cache. + params.slot_mapping = masked_slot_mapping; + if (k_cache_scale.has_value()) { + params.k_cache_scale = k_cache_scale; + xllm::kernel::quant_to_paged_cache(params); + } else { + xllm::kernel::reshape_paged_cache(params); + } + } + + q_input = q_input.view({tokens, tp_heads_.attn, -1}); + auto q_gather_ctx = parallel_state::launch_all_gather( + q_input.contiguous(), dcp_group_); + torch::Tensor q_full = parallel_state::finish_all_gather(std::move(q_gather_ctx)); + const int64_t head_dim = q_input.size(-1); + q_full = q_full.permute({1, 0, 2, 3}).contiguous() + .view({tokens, dcp_size_ * tp_heads_.attn, head_dim}); + + const int64_t local_heads = full_heads_.attn; + torch::Tensor q_decode = q_full.view({tokens, local_heads, head_dim}) + .unsqueeze(1) + .contiguous(); // [tokens, 1, heads, D] + torch::Tensor partial_out = + torch::empty({tokens, 1, local_heads, kv_lora_rank_}, hidden_states.options()); + std::optional partial_lse = + torch::empty({tokens, local_heads, 1}, + hidden_states.options().dtype(torch::kFloat32)); + int64_t kv_cache_quant_bit_size = -1; + std::optional k_cache_quant_scale; + if (k_cache_scale.has_value()) { + k_cache_quant_scale = k_cache_scale; + kv_cache_quant_bit_size = 8; + } + + auto max_seq_len = (local_meta.max_seq_len + dcp_size_ - 1) / dcp_size_; // ceil div + auto kv_seq_lens = (local_meta.kv_seq_lens + dcp_size_ - 1) / dcp_size_; // ceil div + kv_seq_lens = kv_seq_lens.to(torch::kInt32); + xllm::kernel::mlu::batch_decode( + q_decode, + k_cache, + partial_out, + local_meta.block_table, + kv_seq_lens, + /*v_cache=*/std::nullopt, + partial_lse, + /*q_quant_scale=*/std::nullopt, + k_cache_quant_scale, + /*v_cache_quant_scale=*/std::nullopt, + /*out_quant_scale=*/std::nullopt, + /*alibi_slope=*/std::nullopt, + /*mask=*/std::nullopt, + local_meta.compute_dtype, + max_seq_len, + std::max(sliding_window_ - 1, -1), + /*window_size_right=*/-1, + attn_scale_, + /*return_lse=*/true, + kv_cache_quant_bit_size, + /*cu_seq_q=*/std::nullopt, + /*max_seq_q=*/-1, + /*sink=*/std::nullopt); + + torch::Tensor partial_out_flat = + partial_out.permute({2, 0, 1, 3}).contiguous() + .view({local_heads, tokens * kv_lora_rank_}); + auto o_a2a_ctx = parallel_state::launch_all_to_all(partial_out_flat, dcp_group_); + torch::Tensor partial_out_redistributed = + parallel_state::finish_all_to_all(std::move(o_a2a_ctx)); + + torch::Tensor partial_lse_flat = + partial_lse.value().permute({1, 0, 2}).contiguous() + .view({local_heads, tokens}); + auto lse_a2a_ctx = parallel_state::launch_all_to_all(partial_lse_flat, dcp_group_); + torch::Tensor partial_lse_redistributed = + parallel_state::finish_all_to_all(std::move(lse_a2a_ctx)); + + partial_out_redistributed = + partial_out_redistributed.view({dcp_size_, 1, tp_heads_.attn, tokens, kv_lora_rank_}) + .permute({0, 3, 1, 2, 4}).contiguous(); + partial_lse_redistributed = + partial_lse_redistributed.view({dcp_size_, tp_heads_.attn, tokens, 1}) + .permute({0, 2, 1, 3}).contiguous(); + + torch::Tensor merged_out = partial_out_redistributed[0].clone(); + torch::Tensor merged_lse = partial_lse_redistributed[0].clone(); + for (int32_t i = 1; i < dcp_size_; ++i) { + xllm::kernel::mlu::update_out_and_lse( + merged_out, merged_lse, + partial_out_redistributed[i], partial_lse_redistributed[i]); + } + + return project_output(merged_out, tp_heads_); +} + +} // namespace layer +} // namespace xllm diff --git a/xllm/core/platform/platform.h b/xllm/core/platform/platform.h index 906614ce01..a93c465f9e 100644 --- a/xllm/core/platform/platform.h +++ b/xllm/core/platform/platform.h @@ -62,6 +62,7 @@ class Platform final { static constexpr bool supports_dsa_indexer_cache_elision() { return is_mlu(); } + static constexpr bool uses_decode_dcp() { return is_mlu(); } static constexpr bool is_ilu() { #if defined(USE_ILU) diff --git a/xllm/core/runtime/mlu_graph_executor_impl.cpp b/xllm/core/runtime/mlu_graph_executor_impl.cpp index 0c16fa279a..4c3501ace7 100644 --- a/xllm/core/runtime/mlu_graph_executor_impl.cpp +++ b/xllm/core/runtime/mlu_graph_executor_impl.cpp @@ -79,7 +79,8 @@ GraphPoolMemoryUsage get_graph_pool_usage( const auto snapshot = torch_mlu::MLUCachingAllocator::snapshot(); for (const auto& segment : snapshot.segments) { if (segment.device != device_index || - segment.owner_private_pool_id != pool_id) { + segment.owner_private_pool_id.first != static_cast(pool_id.first) || + segment.owner_private_pool_id.second != static_cast(pool_id.second)) { continue; } usage.reserved_bytes += segment.total_size; diff --git a/xllm/core/runtime/options.h b/xllm/core/runtime/options.h index 97c8c7d732..e2fa1f176d 100644 --- a/xllm/core/runtime/options.h +++ b/xllm/core/runtime/options.h @@ -119,6 +119,11 @@ struct Options { // Context parallelism size PROPERTY(int32_t, cp_size) = 1; + // Decode context parallelism size. + PROPERTY(int32_t, dcp_size) = 1; + + PROPERTY(int32_t, cp_kv_cache_interleave_size) = 0; + // tensor parallelism size // Default set as 1 PROPERTY(int32_t, tp_size) = 1; diff --git a/xllm/core/runtime/worker_impl.cpp b/xllm/core/runtime/worker_impl.cpp index 5fa33ae766..c5d3479302 100644 --- a/xllm/core/runtime/worker_impl.cpp +++ b/xllm/core/runtime/worker_impl.cpp @@ -661,7 +661,6 @@ ForwardInput WorkerImpl::update_input_by_last_step_output_for_schedule_overlap( return update_input_by_last_step_output(input); } -#if defined(USE_NPU) torch::Tensor WorkerImpl::recompute_new_cache_slots(const ForwardInput& input) { auto old_cache_slots = input.input_params.attention.device.new_cache_slots; int64_t numel = old_cache_slots.numel(); @@ -700,6 +699,7 @@ torch::Tensor WorkerImpl::recompute_new_cache_slots(const ForwardInput& input) { return new_cache_slots; } +#if defined(USE_NPU) torch::Tensor WorkerImpl::compute_in_prefix_slots(const ForwardInput& input) { // Derive prefix block count from `kv_cache_tokens_nums` (already-cached // tokens at the start of this forward), which covers prefix-cache hits and @@ -779,6 +779,45 @@ void WorkerImpl::prepare_work_before_execute(const ForwardInput& input, input, processed_input, *prepare_stream_); } +torch::Tensor WorkerImpl::recompute_dcp_cache_slots(const ForwardInput& input) { + const int32_t dcp_size = parallel_args_.dcp_size_effective(); + CHECK_GT(dcp_size, 1) << "recompute_dcp_cache_slots requires dcp_size > 1"; + const int32_t dcp_rank = parallel_args_.dcp_rank(); + + torch::Tensor old_cache_slots = + input.input_params.attention.device.new_cache_slots; + if (!old_cache_slots.defined() || old_cache_slots.numel() == 0) { + return old_cache_slots; + } + + const int32_t block_size = options_.block_size(); + torch::Tensor positions = + input.host_positions().to(torch::kCPU).to(torch::kLong); + const int32_t interleave_size = + parallel_args_.cp_kv_cache_interleave_size_effective(block_size); + + const int64_t virtual_block_size = block_size * dcp_size; + torch::Tensor pos = positions.to(torch::kCPU).to(torch::kLong); + torch::Tensor slots = old_cache_slots.to(torch::kCPU).to(torch::kLong); + + torch::Tensor virtual_block_offsets = pos % virtual_block_size; + torch::Tensor owner = (virtual_block_offsets / interleave_size) % dcp_size; + torch::Tensor mask = (owner == dcp_rank); + + torch::Tensor local_offsets = + (virtual_block_offsets / (dcp_size * interleave_size)) * interleave_size + + (virtual_block_offsets % interleave_size); + + // Physical slot: logical_block_id = old_slot / virtual_block_size. + torch::Tensor logical_block_id = slots / virtual_block_size; + torch::Tensor physical_slot = logical_block_id * block_size + local_offsets; + + // Non-owned tokens → -1. + torch::Tensor remapped = torch::where( + mask, physical_slot, torch::full_like(slots, -1, slots.options())); + return remapped.to(old_cache_slots.scalar_type()).to(old_cache_slots.device()); +} + void WorkerImpl::prepare_work_before_execute_on_stream( const ForwardInput& input, ForwardInput& processed_input, @@ -938,6 +977,16 @@ void WorkerImpl::prepare_work_before_execute_on_stream( in_prefix_slots.to(device_); } #endif + const bool needs_dcp_decode_prep = + parallel_args_.dcp_size_effective() > 1 && + processed_input.input_params.meta.batch_forward_type.is_decode() && + !processed_input.cp_partitioned; + if (needs_dcp_decode_prep) { + torch::Tensor new_cache_slots = + recompute_dcp_cache_slots(processed_input); + processed_input.input_params.attention.device.new_cache_slots = + new_cache_slots.to(device_); + } auto& input_params = processed_input.input_params; diff --git a/xllm/core/runtime/worker_impl.h b/xllm/core/runtime/worker_impl.h index a1da04c390..54636b333d 100644 --- a/xllm/core/runtime/worker_impl.h +++ b/xllm/core/runtime/worker_impl.h @@ -275,9 +275,10 @@ class WorkerImpl { // decoder ATB binding refresh. bool init_rolling_runtime_state(); - torch::Tensor recompute_new_cache_slots(const ForwardInput& input); torch::Tensor compute_in_prefix_slots(const ForwardInput& input); #endif + torch::Tensor recompute_new_cache_slots(const ForwardInput& input); + torch::Tensor recompute_dcp_cache_slots(const ForwardInput& input); protected: // runtime options diff --git a/xllm/xllm.cpp b/xllm/xllm.cpp index 24bd4937a5..0889a5952e 100644 --- a/xllm/xllm.cpp +++ b/xllm/xllm.cpp @@ -166,6 +166,7 @@ Options create_options(const std::string& instance_name, bool is_local) { .node_rank(distributed_config.node_rank()) .dp_size(parallel_config.dp_size()) .cp_size(parallel_config.cp_size()) + .dcp_size(parallel_config.dcp_size()) .ep_size(parallel_config.ep_size()) .tp_size(static_cast(parallel_config.tp_size())) .sp_size(static_cast(parallel_config.sp_size()))