diff --git a/ark/api/planner.cpp b/ark/api/planner.cpp index c48e19c5..91c208ac 100644 --- a/ark/api/planner.cpp +++ b/ark/api/planner.cpp @@ -357,6 +357,19 @@ std::string Planner::Impl::plan(bool pretty) const { {"TaskRange", {0, num_tasks}}, {"Granularity", granularity}}}; + // Override NumProcs with the actual processor count from + // ProcessorRange. default_config sets NumProcs = NumTasks, + // but the planner may assign fewer processors. + if (task_infos.back()["Ops"][0].at("Type") == + "AllReducePacketFused") { + auto &pr = resource_group["ProcessorRange"]; + size_t pr_begin = pr[0].get(); + size_t pr_end = pr[1].get(); + size_t pr_step = (pr.size() >= 3) ? pr[2].get() : 1; + size_t num_procs = (pr_end - pr_begin + pr_step - 1) / pr_step; + task_infos.back()["Ops"][0]["Config"]["NumProcs"] = num_procs; + } + if (new_processor_group) { processor_group["ResourceGroups"] = Json::array(); processor_group["ResourceGroups"].push_back(resource_group); diff --git a/ark/include/ark/model.hpp b/ark/include/ark/model.hpp index 8fc083bb..6c1d15cc 100644 --- a/ark/include/ark/model.hpp +++ b/ark/include/ark/model.hpp @@ -271,6 +271,14 @@ class Model : public ModelGraph { // total number of ranks `rank_num`. Tensor all_reduce(Tensor input, int rank, int rank_num, Tensor output = NullTensor, const std::string &name = ""); + // Packet-based all-reduce: uses LL packet channels (data + flag fused) for + // low-latency intra-node collective. Single-shot reduce-scatter + + // allgather, NOT a ring chain — avoids the (N-1)-step task accumulation of + // all_reduce() and is the recommended primitive for small messages on + // NVLink. + Tensor all_reduce_packet(Tensor input, int rank, int rank_num, + Tensor output = NullTensor, + const std::string &name = ""); // Performs an all-gather operator across all ranks, aggregating the input // tensors. Takes the `input` tensor, the current GPU's rank, and the // total number of ranks `rank_num`. Returns a vector of tensors, each diff --git a/ark/include/kernels/comm.h b/ark/include/kernels/comm.h index 4a2deca8..ef44690f 100644 --- a/ark/include/kernels/comm.h +++ b/ark/include/kernels/comm.h @@ -437,6 +437,136 @@ DEVICE void device_sync(int, int) { } } +// Fused intra-node packet allreduce, mirroring mscclpp's `allreducePacket` +// kernel (`src/ext/collectives/allreduce/allreduce_packet.cu`). The three +// phases (scatter → reduce → broadcast) live in a single device function; +// inter-phase synchronization is provided by the LL packet flag itself +// (writer sets data+flag, reader spins on flag), so no cross-block barrier +// is needed. +// +// Template parameters: +// NPeers — world_size - 1; this rank's peer count +// Rank — this rank's world index [0, NPeers] +// NumProcs — number of blocks the planner gave this op +// (= (ProcEnd - ProcBegin) / ProcStep, supplied at codegen) +// NumWarps — warps per block +// PacketType — `mscclpp::LL16Packet` (only LL16 supported initially; +// Payload = uint2 = 4 fp16 elements per packet) +// DataType — element type (only fp16 supported initially) +// NelemsTotal — total elements in the input tensor (compile-time, from +// shape) +// Flag — compile-time flag value used by the LL packet protocol +// (per-iter flag rotation is deferred work) +// +// Runtime arguments: +// output / input / scratch — base pointers +// scratch_offset_remote — offset (in bytes) of the scratch tensor +// within ALL ranks' device buffers; since +// internal-buffer allocation is symmetric +// for symmetric collectives, peers' scratch +// is at the same local offset as ours +// input_offset — offset (in bytes) of the input tensor +// within its device buffer +// task_id — block-relative index in [0, NumProcs) +// +// Block-to-peer mapping requirement (initial implementation): +// NumProcs >= NPeers AND NumProcs % NPeers == 0 +template +DEVICE void allreduce_packet_fused(DataType *output, DataType *input, + void *scratch, + uint64_t scratch_offset_remote, + uint64_t input_offset, int task_id, + int /*sram_per_warp*/) { + using Payload = typename PacketType::Payload; + static_assert(NPeers >= 1, "Need at least one peer"); + static_assert(NumProcs >= NPeers && (NumProcs % NPeers) == 0, + "Initial port requires NumProcs >= NPeers and divisible"); + static_assert(std::is_same::value, + "Only LL16Packet (Payload=uint2) supported initially"); + static_assert(sizeof(DataType) <= sizeof(uint32_t), + "DataType must be <= 4 bytes"); + + constexpr int NThreadsPerBlock = NumWarps * 32; + constexpr int WorldSize = NPeers + 1; + constexpr int ElemsPerUint32 = sizeof(uint32_t) / sizeof(DataType); + static_assert( + NelemsTotal % (ElemsPerUint32 * 2 * WorldSize) == 0, + "NelemsTotal must be divisible by WorldSize * 2 * ElemsPerUint32"); + constexpr int NelemsInt32 = NelemsTotal / ElemsPerUint32; + constexpr int NPkts = NelemsInt32 / 2; + constexpr int NelemsPerRank = NelemsInt32 / WorldSize; + constexpr int NPktsPerRank = NelemsPerRank / 2; + constexpr int NBlocksPerPeer = NumProcs / NPeers; + constexpr int NelemsPerPacket = sizeof(Payload) / sizeof(DataType); + + const int peer_idx = task_id / NBlocksPerPeer; + const int local_block_idx = task_id % NBlocksPerPeer; + const int remote_rank = peer_idx < Rank ? peer_idx : peer_idx + 1; + const int peer_tid = threadIdx.x + local_block_idx * NThreadsPerBlock; + constexpr int peer_total_threads = NThreadsPerBlock * NBlocksPerPeer; + constexpr uint64_t SCRATCH_INPUT_BYTES = NPkts * sizeof(PacketType); + + // ----- Phase 1: putPackets to peer's scratch[Rank slot] ----- + if (task_id < NumProcs) { + auto &chan = ARK_SM_CHANS[remote_rank]; + uint64_t dst_off = + scratch_offset_remote + Rank * NPktsPerRank * sizeof(PacketType); + uint64_t src_off = + input_offset + remote_rank * NelemsPerRank * sizeof(uint32_t); + chan.template putPackets( + dst_off, src_off, NelemsPerRank * sizeof(uint32_t), peer_tid, + peer_total_threads, Flag); + } + + // ----- Phase 2: reduce local rank's shard, scatter result to peers ----- + { + PacketType *scratch_input = reinterpret_cast(scratch); + uint2 *src = reinterpret_cast(input) + Rank * NPktsPerRank; + uint2 *dst = reinterpret_cast(output) + Rank * NPktsPerRank; + + for (int idx = threadIdx.x + task_id * NThreadsPerBlock; + idx < NPktsPerRank; idx += NThreadsPerBlock * NumProcs) { + uint2 data = src[idx]; +#pragma unroll + for (int i = 0; i < NPeers; ++i) { + int peer_rank = i < Rank ? i : i + 1; + PacketType *pkt = + scratch_input + peer_rank * NPktsPerRank + idx; + uint2 val = pkt->read(Flag, -1); + ReduceTypeSum::template reduce( + reinterpret_cast(&data), + reinterpret_cast(&data), + reinterpret_cast(&val)); + } + dst[idx] = data; +#pragma unroll + for (int i = 0; i < NPeers; ++i) { + int peer_rank = i < Rank ? i : i + 1; + auto &chan = ARK_SM_CHANS[peer_rank]; + char *remote_base = reinterpret_cast(chan.dst_); + PacketType *remote_pkt = reinterpret_cast( + remote_base + scratch_offset_remote + SCRATCH_INPUT_BYTES + + Rank * NPktsPerRank * sizeof(PacketType)); + (remote_pkt + idx)->write(data, Flag); + } + } + } + + // ----- Phase 3: read result packets for this peer's shard ----- + if (task_id < NumProcs) { + PacketType *result_base = reinterpret_cast( + reinterpret_cast(scratch) + SCRATCH_INPUT_BYTES + + remote_rank * NPktsPerRank * sizeof(PacketType)); + uint2 *dst = + reinterpret_cast(output) + remote_rank * NPktsPerRank; + for (int idx = peer_tid; idx < NPktsPerRank; + idx += peer_total_threads) { + dst[idx] = result_base[idx].read(Flag, -1); + } + } +} + } // namespace ark #endif // ARK_KERNELS_COMM_H_ diff --git a/ark/model/model_op.cpp b/ark/model/model_op.cpp index 0cf43269..2303d7b5 100644 --- a/ark/model/model_op.cpp +++ b/ark/model/model_op.cpp @@ -94,6 +94,7 @@ const ModelOpType ModelOpT::from_name(const std::string &type_name) { MODEL_OP_TYPE_REGISTER(RecvPacket); MODEL_OP_TYPE_REGISTER(RecvReduceSendPacket); MODEL_OP_TYPE_REGISTER(RecvReduceSend); + MODEL_OP_TYPE_REGISTER(AllReducePacketFused); MODEL_OP_TYPE_REGISTER(DeviceSync); } auto it = instances.find(type_name); diff --git a/ark/ops/ops_all_reduce.cpp b/ark/ops/ops_all_reduce.cpp index 49ee60ac..320194e4 100644 --- a/ark/ops/ops_all_reduce.cpp +++ b/ark/ops/ops_all_reduce.cpp @@ -2,6 +2,7 @@ // Licensed under the MIT license. #include "ops_common.hpp" +#include "ops_communication.hpp" namespace ark { @@ -42,4 +43,47 @@ Tensor Model::all_reduce(Tensor input, int gpu_id, int gpu_num, Tensor output, return cumulate; } +Tensor Model::all_reduce_packet(Tensor input, int rank, int rank_num, + Tensor output, const std::string &) { + int n_peers = rank_num - 1; + if (n_peers < 1) { + ERR(ModelError, "all_reduce_packet requires rank_num >= 2"); + } + + if (output.is_null()) { + output = this->tensor(input.shape(), input.data_type()); + } + + // Scratch layout: [input_section | result_section] + // Each section holds NPkts = nelems_int32 / 2 packets of 16 bytes each. + // Total: 2 × NPkts × 16 = nelems_int32 × 16 = nelems_fp16 × 8 bytes. + size_t nelems = input.shape().nelems(); + size_t elems_per_uint32 = sizeof(uint32_t) / input.data_type().bytes(); + if (nelems % (elems_per_uint32 * 2 * rank_num) != 0) { + ERR(ModelError, "all_reduce_packet: nelems (", nelems, + ") must be divisible by ", elems_per_uint32 * 2 * rank_num); + } + size_t nelems_int32 = nelems / elems_per_uint32; + size_t n_pkts = nelems_int32 / 2; // each packet carries uint2 = 2×u32 + size_t packet_size = 16; // sizeof(mscclpp::LL16Packet) + size_t scratch_bytes = 2 * n_pkts * packet_size; + Dims scratch_shape(static_cast(scratch_bytes)); + Tensor scratch = this->tensor(scratch_shape, UINT8); + + // Peer scratch refs — remote buffers at the same offset (symmetric alloc). + std::vector peer_scratch_refs; + for (int i = 0; i < rank_num; ++i) { + if (i == rank) continue; + peer_scratch_refs.push_back(std::make_shared( + UINT8.ref(), std::make_shared(i), scratch_shape)); + } + + uint32_t flag = 1; // Hardcoded; per-call rotation deferred. + return impl_ + ->create_op( + "all_reduce_packet", input.ref(), output.ref(), rank, rank_num, + flag, scratch.ref(), peer_scratch_refs) + ->result_tensors()[0]; +} + } // namespace ark diff --git a/ark/ops/ops_communication.cpp b/ark/ops/ops_communication.cpp index 4e221e17..b5dd4061 100644 --- a/ark/ops/ops_communication.cpp +++ b/ark/ops/ops_communication.cpp @@ -452,6 +452,100 @@ Json ModelOpRecvReduceSendPacket::default_config( return config; } +// Fused intra-node packet allreduce. Replaces the three-op chain +// (send_packet → recv_reduce_send_packet → recv_packet) with a single op. +ModelOpAllReducePacketFused::ModelOpAllReducePacketFused( + ModelTensorRef input, ModelTensorRef output, int rank, int rank_num, + uint32_t flag, ModelTensorRef scratch, + const std::vector &peer_scratch_refs) + : ModelOp("AllReducePacketFused") { + check_null(input); + check_null(output); + check_null(scratch); + if (scratch->buffer()->rank() != rank && scratch->buffer()->rank() != -1) { + ERR(ModelError, + "invalid local scratch buffer rank: ", scratch->buffer()->rank(), + ", expected: ", rank); + } + uint32_t n_peers = rank_num - 1; + if (peer_scratch_refs.size() != n_peers) { + ERR(ModelError, "expected ", n_peers, " peer scratch refs, got ", + peer_scratch_refs.size()); + } + // The peer scratches are listed as write tensors purely so that + // `CommResource::connect()` discovers the remote ranks and sets up SM + // channels to them. The kernel uses `ARK_SM_CHANS[remote_rank].dst_` + // directly + the local scratch offset (which is identical on every rank + // for symmetric internal-buffer allocation). + read_tensors_ = {input, scratch}; + write_tensors_ = {output}; + for (auto &p : peer_scratch_refs) { + write_tensors_.push_back(p); + } + ModelTensorRef result = std::make_shared(*output); + result_tensors_ = {result}; + args_ = { + {"Flag", ModelOpArg(flag)}, + {"NPeers", ModelOpArg(n_peers)}, + {"Rank", ModelOpArg(rank)}, + }; + verify(); +} + +std::string ModelOpAllReducePacketFused::impl_name(const Json &config) const { + check_fields_config(config, {"NumProcs", "NumWarps", "PacketType", + "NumTasks", "SramBytes", "Tile"}); + auto &input = read_tensors_[0]; + uint32_t flag = args_.at("Flag").value(); + uint32_t n_peers = args_.at("NPeers").value(); + int rank = args_.at("Rank").value(); + int num_warps = config.at("NumWarps"); + int num_procs = config.at("NumProcs"); + std::string packet_type = config.at("PacketType"); + // Total fp16 elements in input — compile-time constant. + DimType nelems_total = input->shape().nelems(); + return function_name_string( + "allreduce_packet_fused", + {std::to_string(n_peers), std::to_string(rank), + std::to_string(num_procs), std::to_string(num_warps), packet_type, + input->data_type()->type_str(), std::to_string(nelems_total), + std::to_string(flag)}); +} + +std::vector ModelOpAllReducePacketFused::impl_args( + [[maybe_unused]] const Json &config) const { + // (output, input, scratch_ptr, scratch_offset_remote, input_offset) + return {write_tensors_[0], read_tensors_[0], read_tensors_[1], + ModelOffset(read_tensors_[1]), ModelOffset(read_tensors_[0])}; +} + +Json ModelOpAllReducePacketFused::default_config( + [[maybe_unused]] const ArchRef arch) const { + Json config; + config["PacketType"] = "mscclpp::LL16Packet"; + // Tuning follows mscclpp's `getDefaultBlockNumAndThreadNum` in + // src/ext/collectives/allreduce/allreduce_packet.cu: + // - small (<32 KB): 4 blocks per peer, 1024 threads + // - larger: 8 blocks per peer, 512 threads (<=153KB) or 1024 + // We translate "blocks per peer × n_peers" into NumTasks (the planner + // hands us that many blocks, capped by the ProcessorRange when the + // user supplies a PlannerContext). + auto &input = read_tensors_[0]; + DimType input_bytes = input->shape().nelems() * input->data_type()->bytes(); + uint32_t n_peers = args_.at("NPeers").value(); + int blocks_per_peer = (input_bytes < (32 << 10)) ? 4 : 8; + int num_warps = 32; // 1024 threads / 32 = 32 warps per block + if (input_bytes >= (32 << 10) && input_bytes <= 153600) { + num_warps = 16; // 512 threads + } + config["NumWarps"] = num_warps; + config["SramBytes"] = 0; + config["Tile"] = {1, 1}; // not used by this op, but required by codegen + config["NumTasks"] = blocks_per_peer * static_cast(n_peers); + config["NumProcs"] = config["NumTasks"]; + return config; +} + ModelOpRecvReduceSend::ModelOpRecvReduceSend( ModelTensorRef input, ModelTensorRef output, int rank, const std::vector &remote_ranks, int recv_tag, int output_tag, diff --git a/ark/ops/ops_communication.hpp b/ark/ops/ops_communication.hpp index f0c0134f..edd20b5f 100644 --- a/ark/ops/ops_communication.hpp +++ b/ark/ops/ops_communication.hpp @@ -103,6 +103,26 @@ class ModelOpRecvReduceSend : public ModelOp { Json default_config(const ArchRef arch = ARCH_ANY) const override; }; +// Fused intra-node packet allreduce — one device function emits the three +// phases of the LL-packet allreduce (scatter → reduce → broadcast). +// Replaces the three-op chain (`send_packet`, +// `recv_reduce_send_packet`, `recv_packet`) used by the old +// `Model::all_reduce_packet`. +class ModelOpAllReducePacketFused : public ModelOp { + public: + ModelOpAllReducePacketFused() = default; + ModelOpAllReducePacketFused( + ModelTensorRef input, ModelTensorRef output, int rank, int rank_num, + uint32_t flag, ModelTensorRef scratch, + const std::vector &peer_scratch_refs); + + std::string impl_name(const Json &config) const override; + + std::vector impl_args(const Json &config) const override; + + Json default_config(const ArchRef arch = ARCH_ANY) const override; +}; + class ModelOpDeviceSync : public ModelOp { public: ModelOpDeviceSync() = default; diff --git a/ark/ops/ops_communication_test.cpp b/ark/ops/ops_communication_test.cpp index 5c21d6fe..680ee32f 100644 --- a/ark/ops/ops_communication_test.cpp +++ b/ark/ops/ops_communication_test.cpp @@ -490,10 +490,17 @@ ark::unittest::State test_communication_send_recv_reduce() { auto op = nlohmann::json::parse(op_str); nlohmann::json config; if (op.at("Type") == "Send") { + constexpr int tile_y = 256; + const auto &shape = op.at("WriteTensors")[0].at("PaddedShape"); + size_t num_tasks = 1; + for (const auto &dim : shape) { + num_tasks *= dim.get(); + } + num_tasks = (num_tasks + tile_y - 1) / tile_y; config["ChannelType"] = "Sm"; config["Signal"] = false; - config["Tile"] = {1, 256}; - config["NumTasks"] = 2; + config["Tile"] = {1, tile_y}; + config["NumTasks"] = num_tasks; config["NumWarps"] = 4; config["SramBytes"] = 0; } else if (op.at("Type") == "DeviceSync") { @@ -558,6 +565,136 @@ ark::unittest::State test_communication_send_recv_reduce() { return ark::unittest::SUCCESS; } +ark::unittest::State test_communication_allreduce_packet_fused_model() { + // Single-GPU model-level test: construct the fused allreduce op and + // verify impl_name / impl_args / default_config produce valid output. + { + ark::Model model(0, 2); + ark::Tensor tns = model.tensor({1024}, ark::FP16); + ark::Tensor result = model.all_reduce_packet(tns, 0, 2); + + auto nodes = model.nodes(); + bool found = false; + for (auto &node : nodes) { + auto &op = node->op; + if (op->is_virtual()) continue; + if (op->type() != ark::ModelOpT::from_name("AllReducePacketFused")) + continue; + found = true; + auto cfg = op->default_config(ark::ARCH_CUDA_80); + UNITTEST_FALSE(cfg.empty()); + auto name = op->impl_name(cfg); + UNITTEST_FALSE(name.empty()); + // Verify the kernel name appears in the impl_name string. + UNITTEST_TRUE(name.find("allreduce_packet_fused") != + std::string::npos); + auto args = op->impl_args(cfg); + // (output, input, scratch_ptr, scratch_offset_remote, input_offset) + UNITTEST_EQ(args.size(), 5); + } + UNITTEST_TRUE(found); + } + // Medium-size tensor (40000 bytes, 32KB <= x <= 153KB): + // exercises the blocks_per_peer=8, num_warps=16 config path. + { + // 20000 FP16 = 40000 bytes => 32KB < 40KB <= 153KB. + ark::Model model(0, 2); + ark::Tensor tns = model.tensor({20000}, ark::FP16); + ark::Tensor result = model.all_reduce_packet(tns, 0, 2); + + auto nodes = model.nodes(); + for (auto &node : nodes) { + auto &op = node->op; + if (op->is_virtual()) continue; + if (op->type() != ark::ModelOpT::from_name("AllReducePacketFused")) + continue; + auto cfg = op->default_config(ark::ARCH_CUDA_80); + UNITTEST_EQ(cfg.at("NumWarps").get(), 16); + UNITTEST_EQ(cfg.at("NumTasks").get(), 8); + } + } + // Large tensor (> 153600 bytes): + // exercises blocks_per_peer=8, num_warps=32 config path. + { + // 80000 FP16 = 160000 bytes > 153600. + ark::Model model(0, 2); + ark::Tensor tns = model.tensor({80000}, ark::FP16); + ark::Tensor result = model.all_reduce_packet(tns, 0, 2); + + auto nodes = model.nodes(); + for (auto &node : nodes) { + auto &op = node->op; + if (op->is_virtual()) continue; + if (op->type() != ark::ModelOpT::from_name("AllReducePacketFused")) + continue; + auto cfg = op->default_config(ark::ARCH_CUDA_80); + UNITTEST_EQ(cfg.at("NumWarps").get(), 32); + UNITTEST_EQ(cfg.at("NumTasks").get(), 8); + } + } + // Test Planner::plan() with AllReducePacketFused to exercise the + // NumProcs override branch in planner.cpp. + { + ark::Model model(0, 2); + ark::Tensor tns = model.tensor({1024}, ark::FP16); + ark::Tensor result = model.all_reduce_packet(tns, 0, 2); + + ark::Planner planner(model, 0); + auto plan = ark::Json::parse(planner.plan(false)); + // The plan should contain at least one TaskInfo for the fused op. + UNITTEST_TRUE(plan.contains("TaskInfos")); + bool found_fused = false; + for (auto &ti : plan["TaskInfos"]) { + for (auto &op : ti["Ops"]) { + if (op.at("Type").get() == + "AllReducePacketFused") { + found_fused = true; + // NumProcs should have been set by the planner. + UNITTEST_TRUE(op["Config"].contains("NumProcs")); + } + } + } + UNITTEST_TRUE(found_fused); + } + // Verify rank_num < 2 is rejected. + { + ark::Model model(0, 1); + ark::Tensor tns = model.tensor({1024}, ark::FP16); + UNITTEST_THROW(model.all_reduce_packet(tns, 0, 1), ark::ModelError); + } + // Verify non-divisible tensor size is rejected. + { + ark::Model model(0, 2); + ark::Tensor tns = model.tensor({1023}, ark::FP16); + UNITTEST_THROW(model.all_reduce_packet(tns, 0, 2), ark::ModelError); + } + // Multi-peer model test: rank_num=4, rank=2 exercises peer-index mapping. + { + ark::Model model(2, 4); + ark::Tensor tns = model.tensor({1024}, ark::FP16); + ark::Tensor result = model.all_reduce_packet(tns, 2, 4); + + auto nodes = model.nodes(); + bool found = false; + for (auto &node : nodes) { + auto &op = node->op; + if (op->is_virtual()) continue; + if (op->type() != ark::ModelOpT::from_name("AllReducePacketFused")) + continue; + found = true; + auto cfg = op->default_config(ark::ARCH_CUDA_80); + // 1024 FP16 = 2048 bytes < 32KB → blocks_per_peer=4, n_peers=3 + UNITTEST_EQ(cfg.at("NumTasks").get(), 12); + UNITTEST_EQ(cfg.at("NumWarps").get(), 32); + auto name = op->impl_name(cfg); + UNITTEST_TRUE(name.find("allreduce_packet_fused") != + std::string::npos); + } + UNITTEST_TRUE(found); + } + return ark::unittest::SUCCESS; +} + int main() { ark::init(); UNITTEST(test_communication_host_ops); @@ -567,5 +704,6 @@ int main() { UNITTEST(test_communication_send_packet); UNITTEST(test_communication_send_recv_reduce_packet); UNITTEST(test_communication_send_recv_reduce); + UNITTEST(test_communication_allreduce_packet_fused_model); return ark::unittest::SUCCESS; } diff --git a/python/ark/ops.py b/python/ark/ops.py index 2b7e387f..477bc817 100644 --- a/python/ark/ops.py +++ b/python/ark/ops.py @@ -36,6 +36,7 @@ "mul", "div", "all_reduce", + "all_reduce_packet", "embedding", "cast", "copy", @@ -598,6 +599,40 @@ def all_reduce( return Tensor(_tensor) +def all_reduce_packet( + input: Tensor, + rank: int, + world_size: int, + output: Tensor = NullTensor, + name: str = "", +) -> Tensor: + """ + Packet-based intra-node all-reduce — single-shot reduce-scatter + allgather using + LL packet channels (fused data+flag write). NOT a ring chain; per-rank + cost is constant in `world_size` for the in-block portion, and the entire + operation completes in 3 task phases regardless of `world_size`. + + Args: + input (Tensor): The input tensor to be reduced. Must be contiguous. + rank (int): The rank of the current process. + world_size (int): The total number of processes. The tensor's + element count must be divisible by (4 * world_size) for FP16. + output (Tensor, optional): The output tensor. If provided, the result + will be stored in this tensor. Defaults to NullTensor. + name (str, optional): The name of the operation. Defaults to + "". + + Returns: + Tensor: The reduced tensor. + """ + if output is not NullTensor: + output = output._tensor + _tensor = Model.get_model().all_reduce_packet( + input._tensor, rank, world_size, output, name + ) + return Tensor(_tensor) + + # def im2col( # input: Tensor, # kernel_height: int, diff --git a/python/model_py.cpp b/python/model_py.cpp index 30b16446..61fc8166 100644 --- a/python/model_py.cpp +++ b/python/model_py.cpp @@ -137,5 +137,8 @@ void register_model(py::module &m) { py::arg("remote_rank"), py::arg("tag"), py::arg("name")) .def("all_reduce", &ark::Model::all_reduce, py::arg("input"), py::arg("rank"), py::arg("world_size"), py::arg("output"), - py::arg("name")); + py::arg("name")) + .def("all_reduce_packet", &ark::Model::all_reduce_packet, + py::arg("input"), py::arg("rank"), py::arg("world_size"), + py::arg("output"), py::arg("name")); } diff --git a/python/unittest/test_ops.py b/python/unittest/test_ops.py index 3e8ee562..b34c7c6c 100644 --- a/python/unittest/test_ops.py +++ b/python/unittest/test_ops.py @@ -216,3 +216,11 @@ def test_ops_transpose(): a = ark.tensor([64, 32], ark.fp16) b = ark.transpose(a, [1, 0]) assert b.shape() == [32, 64] + + +@pytest_ark() +def test_ops_all_reduce_packet(): + ark.set_world_size(2) + a = ark.tensor([1024], ark.fp16) + b = ark.all_reduce_packet(a, rank=0, world_size=2) + assert b.shape() == [1024]