Add monolithic fused packet allreduce kernel and Model::all_reduce_packet API

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>();
+                size_t pr_end = pr[1].get<size_t>();
+                size_t pr_step = (pr.size() >= 3) ? pr[2].get<size_t>() : 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);

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

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 <int NPeers, int Rank, int NumProcs, int NumWarps, typename PacketType,
+          typename DataType, int NelemsTotal, uint32_t Flag = 1>
+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<Payload, uint2>::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<PacketType>(
+            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<PacketType *>(scratch);
+        uint2 *src = reinterpret_cast<uint2 *>(input) + Rank * NPktsPerRank;
+        uint2 *dst = reinterpret_cast<uint2 *>(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<NelemsPerPacket, DataType>(
+                    reinterpret_cast<DataType *>(&data),
+                    reinterpret_cast<DataType *>(&data),
+                    reinterpret_cast<DataType *>(&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<char *>(chan.dst_);
+                PacketType *remote_pkt = reinterpret_cast<PacketType *>(
+                    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<PacketType *>(
+            reinterpret_cast<char *>(scratch) + SCRATCH_INPUT_BYTES +
+            remote_rank * NPktsPerRank * sizeof(PacketType));
+        uint2 *dst =
+            reinterpret_cast<uint2 *>(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_

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);

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<DimType>(scratch_bytes));
+    Tensor scratch = this->tensor(scratch_shape, UINT8);
+
+    // Peer scratch refs — remote buffers at the same offset (symmetric alloc).
+    std::vector<ModelTensorRef> peer_scratch_refs;
+    for (int i = 0; i < rank_num; ++i) {
+        if (i == rank) continue;
+        peer_scratch_refs.push_back(std::make_shared<ModelTensor>(
+            UINT8.ref(), std::make_shared<ModelBuffer>(i), scratch_shape));
+    }
+
+    uint32_t flag = 1;  // Hardcoded; per-call rotation deferred.
+    return impl_
+        ->create_op<ModelOpAllReducePacketFused>(
+            "all_reduce_packet", input.ref(), output.ref(), rank, rank_num,
+            flag, scratch.ref(), peer_scratch_refs)
+        ->result_tensors()[0];
+}
+
 }  // namespace ark

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<ModelTensorRef> &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<ModelTensor>(*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>();
+    uint32_t n_peers = args_.at("NPeers").value<uint32_t>();
+    int rank = args_.at("Rank").value<int>();
+    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<ModelOpArg> 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<uint32_t>();
+    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<int>(n_peers);
+    config["NumProcs"] = config["NumTasks"];
+    return config;
+}
+
 ModelOpRecvReduceSend::ModelOpRecvReduceSend(
     ModelTensorRef input, ModelTensorRef output, int rank,
     const std::vector<int> &remote_ranks, int recv_tag, int output_tag,

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<ModelTensorRef> &peer_scratch_refs);
+
+    std::string impl_name(const Json &config) const override;
+
+    std::vector<ModelOpArg> impl_args(const Json &config) const override;
+
+    Json default_config(const ArchRef arch = ARCH_ANY) const override;
+};
+
 class ModelOpDeviceSync : public ModelOp {
    public:
     ModelOpDeviceSync() = default;