examples/qwen3: TP all-reduce via mscclpp fused-packet, with equivalence test and microbenchmark

.codecov.yml

@@ -22,7 +22,7 @@ ignore:
   - "*/dist-packages/*"
   - "*/third_party/*"
   - "*/ark/*_test.*"
-  - "*/examples/*"
+  - "examples/**"
   - "*/python/unittest/*"
   - "*/ark/unittest/*"
   - "*/ark/ops/ops_test_common.*"

.github/workflows/ut.yml

@@ -87,6 +87,16 @@ jobs:
               --verbose \
               ../python/unittest/
 
+      - name: Run Qwen3 Example Tests
+        if: github.event_name != 'schedule'
+        run: |
+          cd build
+          PYTHONPATH=$PWD/python ARK_ROOT=$PWD python3 -m pytest \
+              --cov=../examples/qwen3 \
+              --cov-report lcov:qwen3_coverage.info \
+              --verbose \
+              ../examples/qwen3/
+
       - name: C++ Coverage
         if: github.event_name != 'schedule'
         run: |
@@ -111,7 +121,7 @@ jobs:
           CODECOV_TOKEN: ${{ secrets.CODECOV_TOKEN }}
         run: |
           cd build
-          lcov -a cpp_coverage.info -a py_coverage.info -o coverage.info
+          lcov -a cpp_coverage.info -a py_coverage.info $([ -f qwen3_coverage.info ] && echo "-a qwen3_coverage.info") -o coverage.info
           bash <(curl -s https://codecov.io/bash) -f coverage.info || echo "Codecov did not collect coverage reports"
 
       - name: Install Python

ark/codegen.cpp

@@ -315,16 +315,17 @@ std::string CodeGenerator::Impl::def_task(const Json &task_json) {
                 size_t buffer_id = moff.buffer_id();
                 auto buf_info = buf_reg.get(buffer_id);
                 if (buf_info && buf_info->is_external) {
-                    ERR(InternalError, "cannot offset external buffer");
-                }
-                size_t buffer_offset;
-                auto it = buffer_id_to_offset_.find(buffer_id);
-                if (it == buffer_id_to_offset_.end()) {
-                    ERR(InternalError, "buffer ID not found: ", buffer_id);
+                    // External buffer: offset relative to its own base.
+                    ss_desc << moff.value();
+                } else {
+                    auto it = buffer_id_to_offset_.find(buffer_id);
+                    if (it == buffer_id_to_offset_.end()) {
+                        ERR(InternalError, "buffer ID not found: ", buffer_id);
+                    }
+                    size_t buffer_offset = it->second;
+                    size_t offset = buffer_offset + moff.value();
+                    ss_desc << offset;
                 }
-                buffer_offset = it->second;
-                size_t offset = buffer_offset + moff.value();
-                ss_desc << offset;
             } else {
                 ss_desc << arg.serialize().begin().value();
             }

ark/codegen_test.cpp

@@ -0,0 +1,173 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+
+#include "codegen.hpp"
+
+#include <map>
+#include <nlohmann/json.hpp>
+#include <set>
+
+#include "ark/model.hpp"
+#include "ark/planner.hpp"
+#include "buffer_registry.hpp"
+#include "model/model_buffer.hpp"
+#include "model/model_node.hpp"
+#include "model/model_op.hpp"
+#include "model/model_op_arg.hpp"
+#include "model/model_tensor.hpp"
+#include "unittest/unittest_utils.h"
+
+// Collect all buffer IDs referenced by OFFSET args in a plan's TaskInfos.
+static std::set<size_t> collect_offset_buffer_ids(const ark::Json &plan) {
+    std::set<size_t> ids;
+    for (auto &ti : plan.at("TaskInfos")) {
+        for (auto &op_json : ti.at("Ops")) {
+            auto op = ark::ModelOp::deserialize(op_json);
+            auto args = op->impl_args(op_json.at("Config"));
+            for (auto &arg : args) {
+                if (arg.type_name() == "OFFSET") {
+                    ids.insert(arg.value<ark::ModelOffset>().buffer_id());
+                }
+            }
+        }
+    }
+    return ids;
+}
+
+// Collect all buffer IDs referenced by TENSOR args in a plan's TaskInfos.
+static std::set<size_t> collect_tensor_buffer_ids(const ark::Json &plan) {
+    std::set<size_t> ids;
+    for (auto &ti : plan.at("TaskInfos")) {
+        for (auto &op_json : ti.at("Ops")) {
+            auto op = ark::ModelOp::deserialize(op_json);
+            auto args = op->impl_args(op_json.at("Config"));
+            for (auto &arg : args) {
+                if (arg.type_name() == "TENSOR") {
+                    ids.insert(
+                        arg.value<ark::ModelTensorRef>()->buffer()->id());
+                }
+            }
+        }
+    }
+    return ids;
+}
+
+// Test 1: CodeGenerator exercises the external-buffer OFFSET path
+// (codegen.cpp line 319: `ss_desc << moff.value();`).
+ark::unittest::State test_codegen_external_buffer_offset() {
+    // Build a 2-rank model with a send_packet op on rank 0.
+    // send_packet's impl_args are two OFFSET args whose buffer IDs we will
+    // register as external in BufferRegistry before constructing CodeGenerator.
+    ark::Model model(0, 2);
+    ark::Tensor tns = model.tensor({1024}, ark::FP16);
+    model.send_packet(tns, 1, /*tag=*/0, /*flag=*/1);
+
+    // Plan on GPU 0.
+    ark::Planner planner(model, 0);
+    auto plan = ark::Json::parse(planner.plan(false));
+
+    // Verify the plan has TaskInfos.
+    UNITTEST_TRUE(plan.contains("TaskInfos"));
+    UNITTEST_TRUE(plan["TaskInfos"].size() > 0);
+
+    // Collect OFFSET and TENSOR buffer IDs from the plan.
+    auto offset_buf_ids = collect_offset_buffer_ids(plan);
+    auto tensor_buf_ids = collect_tensor_buffer_ids(plan);
+    UNITTEST_TRUE(offset_buf_ids.size() > 0);
+
+    // Register every OFFSET buffer as external in BufferRegistry.
+    // Use a dummy non-null pointer; CodeGenerator only checks is_external,
+    // it does not dereference the pointer.
+    auto &buf_reg = ark::BufferRegistry::get_instance();
+    for (size_t id : offset_buf_ids) {
+        buf_reg.set(id, reinterpret_cast<void *>(0x1), 0, /*is_external=*/true);
+    }
+
+    // All referenced buffer IDs go into extra_buffer_ids (external).
+    std::set<size_t> extra;
+    extra.insert(offset_buf_ids.begin(), offset_buf_ids.end());
+    extra.insert(tensor_buf_ids.begin(), tensor_buf_ids.end());
+
+    // Construct CodeGenerator — exercises the external OFFSET path.
+    ark::PlanJson pj(plan);
+    ark::CodeGenerator codegen(pj, /*buffer_id_to_offset=*/{}, extra);
+
+    // Verify non-empty generated code.
+    std::string code = codegen.code();
+    UNITTEST_TRUE(code.size() > 0);
+
+    return ark::unittest::State::SUCCESS;
+}
+
+// Test 2: CodeGenerator exercises the normal (non-external) OFFSET path
+// (codegen.cpp lines 320-325: buffer_id_to_offset_ lookup).
+// Also exercises Model::all_reduce_packet which covers the new
+// `input = this->copy(input)` line in ops_all_reduce.cpp:57.
+ark::unittest::State test_codegen_normal_offset() {
+    // Build a 2-rank model using all_reduce_packet (exercises
+    // ops_all_reduce.cpp line 57: `input = this->copy(input)`).
+    ark::Model model(0, 2);
+    ark::Tensor tns = model.tensor({1024}, ark::FP16);
+    model.all_reduce_packet(tns, 0, 2);
+
+    // Plan on GPU 0.
+    ark::Planner planner(model, 0);
+    auto plan = ark::Json::parse(planner.plan(false));
+    UNITTEST_TRUE(plan.contains("TaskInfos"));
+    UNITTEST_TRUE(plan["TaskInfos"].size() > 0);
+
+    // Collect ALL buffer IDs (OFFSET + TENSOR) from the plan.
+    auto offset_buf_ids = collect_offset_buffer_ids(plan);
+    auto tensor_buf_ids = collect_tensor_buffer_ids(plan);
+    UNITTEST_TRUE(offset_buf_ids.size() > 0);
+
+    // Put all buffer IDs in buffer_id_to_offset_ with offset 0.
+    // Do NOT register them as external in BufferRegistry.
+    std::map<size_t, size_t> buf_id_to_offset;
+    for (size_t id : offset_buf_ids) {
+        buf_id_to_offset[id] = 0;
+    }
+    for (size_t id : tensor_buf_ids) {
+        buf_id_to_offset[id] = 0;
+    }
+
+    // Construct CodeGenerator — exercises the normal OFFSET path
+    // (buffer_id_to_offset_ lookup, lines 320-325 of codegen.cpp).
+    ark::PlanJson pj(plan);
+    ark::CodeGenerator codegen(pj, buf_id_to_offset, {});
+
+    std::string code = codegen.code();
+    UNITTEST_TRUE(code.size() > 0);
+
+    return ark::unittest::State::SUCCESS;
+}
+
+// Test 3: CodeGenerator throws InternalError when an OFFSET arg's buffer ID
+// is neither external nor in buffer_id_to_offset (codegen.cpp line 323).
+ark::unittest::State test_codegen_missing_buffer_id() {
+    // Build a fresh model so its buffer IDs are new (not in BufferRegistry
+    // from test 1).
+    // Safe because ModelBuffer::curr_id is a monotonically-increasing static
+    // counter — IDs allocated here will never collide with test 1's entries.
+    ark::Model model(0, 2);
+    ark::Tensor tns = model.tensor({512}, ark::FP16);
+    model.send_packet(tns, 1, /*tag=*/0, /*flag=*/1);
+
+    ark::Planner planner(model, 0);
+    auto plan = ark::Json::parse(planner.plan(false));
+
+    // Do NOT register any buffer in BufferRegistry.
+    // Do NOT populate buffer_id_to_offset.
+    // CodeGenerator should throw InternalError on the OFFSET lookup.
+    ark::PlanJson pj(plan);
+    UNITTEST_THROW(ark::CodeGenerator(pj, {}, {}), ark::InternalError);
+
+    return ark::unittest::State::SUCCESS;
+}
+
+int main() {
+    UNITTEST(test_codegen_external_buffer_offset);
+    UNITTEST(test_codegen_normal_offset);
+    UNITTEST(test_codegen_missing_buffer_id);
+    return 0;
+}

ark/ops/ops_all_reduce.cpp

@@ -50,6 +50,12 @@ Tensor Model::all_reduce_packet(Tensor input, int rank, int rank_num,
         ERR(ModelError, "all_reduce_packet requires rank_num >= 2");
     }
 
+    // Copy input into an internal buffer so it resides in mscclpp
+    // registered memory.  External buffers (e.g. from torch tensors)
+    // are not part of the registered allocation, and putPackets needs
+    // to read from a registered source.
+    input = this->copy(input);
+
     if (output.is_null()) {
         output = this->tensor(input.shape(), input.data_type());
     }

ark/ops/ops_all_reduce_test.cpp

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
 
+#include "ark/executor.hpp"
+#include "gpu/gpu.hpp"
 #include "model/model_buffer.hpp"
 #include "model/model_node.hpp"
 #include "model/model_op.hpp"
@@ -334,11 +336,109 @@ ark::unittest::State test_all_reduce_sm_8gpus() {
     return ark::unittest::SUCCESS;
 }
 
+template <int NumGpus>
+void test_all_reduce_packet_fused_internal(ark::DimType nelem) {
+    for (int gpu_id = 0; gpu_id < NumGpus; ++gpu_id) {
+        ark::unittest::spawn_process([gpu_id, nelem]() {
+            UNITTEST_SKIP(ark::unittest::get_gpu_count() < NumGpus);
+            // Each GPU's data is equal to its GPU ID + 1.
+            ark::Model m(gpu_id, NumGpus);
+            ark::Tensor ones = m.tensor({nelem}, ark::FP16);
+            ark::Tensor data = m.mul(ones, float(gpu_id + 1));
+            ark::Tensor output = m.all_reduce_packet(data, gpu_id, NumGpus);
+
+            std::vector<ark::half_t> ones_vec(ones.shape().nelems(),
+                                              ark::half_t(1.0f));
+            auto result = ark::op_test(
+                "all_reduce_packet_fused", m, {ones}, {output},
+                baseline_all_reduce<ark::half_t, NumGpus>, {ones_vec.data()});
+            UNITTEST_LOG(result);
+            UNITTEST_EQ(result.max_diff[0], 0.0f);
+            return ark::unittest::SUCCESS;
+        });
+    }
+    ark::unittest::wait_all_processes();
+}
+
+// Variant with external-buffer (placeholder) input — exercises the
+// codegen external-buffer OFFSET path added for all_reduce_packet's
+// internal copy.  Cannot use op_test() because placeholders require
+// pre-allocated GPU memory; drive the executor manually instead.
+template <int NumGpus>
+void test_all_reduce_packet_fused_ext_internal(ark::DimType nelem) {
+    for (int gpu_id = 0; gpu_id < NumGpus; ++gpu_id) {
+        ark::unittest::spawn_process([gpu_id, nelem]() {
+            UNITTEST_SKIP(ark::unittest::get_gpu_count() < NumGpus);
+
+            UNITTEST_EQ(ark::gpuSetDevice(gpu_id), ark::gpuSuccess);
+
+            // Allocate GPU memory and fill with (gpu_id + 1).
+            ark::half_t *d_input = nullptr;
+            size_t nbytes = nelem * sizeof(ark::half_t);
+            UNITTEST_EQ(ark::gpuMalloc(&d_input, nbytes), ark::gpuSuccess);
+            std::vector<ark::half_t> h_input(nelem,
+                                             ark::half_t(float(gpu_id + 1)));
+            UNITTEST_EQ(ark::gpuMemcpy(d_input, h_input.data(), nbytes,
+                                       ark::gpuMemcpyHostToDevice),
+                        ark::gpuSuccess);
+
+            ark::Model m(gpu_id, NumGpus);
+            ark::Tensor input =
+                m.placeholder({nelem}, ark::FP16, {}, {}, {}, -1, d_input);
+            ark::Tensor output = m.all_reduce_packet(input, gpu_id, NumGpus);
+
+            ark::DefaultExecutor exe(m, gpu_id);
+            exe.launch();
+            exe.run(1);
+            exe.stop();
+
+            std::vector<ark::half_t> h_output(nelem);
+            exe.tensor_read(output, h_output);
+
+            float expected = float(NumGpus * (NumGpus + 1)) / 2.0f;
+            for (ark::DimType i = 0; i < nelem; ++i) {
+                UNITTEST_EQ(float(h_output[i]), expected);
+            }
+
+            UNITTEST_EQ(ark::gpuFree(d_input), ark::gpuSuccess);
+            return ark::unittest::SUCCESS;
+        });
+    }
+    ark::unittest::wait_all_processes();
+}
+
+ark::unittest::State test_all_reduce_packet_fused_ext_2gpus() {
+    test_all_reduce_packet_fused_ext_internal<2>(4096);
+    return ark::unittest::SUCCESS;
+}
+
+ark::unittest::State test_all_reduce_packet_fused_2gpus() {
+    test_all_reduce_packet_fused_internal<2>(4096);
+    test_all_reduce_packet_fused_internal<2>(8192);
+    return ark::unittest::SUCCESS;
+}
+
+ark::unittest::State test_all_reduce_packet_fused_4gpus() {
+    test_all_reduce_packet_fused_internal<4>(2048);
+    test_all_reduce_packet_fused_internal<4>(8192);
+    return ark::unittest::SUCCESS;
+}
+
+ark::unittest::State test_all_reduce_packet_fused_8gpus() {
+    test_all_reduce_packet_fused_internal<8>(2048);
+    test_all_reduce_packet_fused_internal<8>(8192);
+    return ark::unittest::SUCCESS;
+}
+
 int main() {
     UNITTEST(test_all_reduce_4gpus);
     UNITTEST(test_all_reduce_8gpus);
     UNITTEST(test_all_reduce_packet_4gpus);
     UNITTEST(test_all_reduce_packet_8gpus);
+    UNITTEST(test_all_reduce_packet_fused_ext_2gpus);
+    UNITTEST(test_all_reduce_packet_fused_2gpus);
+    UNITTEST(test_all_reduce_packet_fused_4gpus);
+    UNITTEST(test_all_reduce_packet_fused_8gpus);
     UNITTEST(test_all_reduce_sm_4gpus);
     UNITTEST(test_all_reduce_sm_8gpus);
     UNITTEST(test_all_reduce_inplace_2gpus);

examples/qwen3/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT license.
+
+# Submodules:
+#   ark_allreduce  - all-reduce wrapper (active)
+#   test_allreduce - tests (active)
+#   bench_allreduce - benchmarks (active)
+#   equiv          - equivalence test helper (staged for inference pipeline)
+#   microbench     - CUDA microbench harness (staged for profiling scripts)
+#   qwen3_config   - model config dataclass (staged for inference pipeline)