Merge remote-tracking branch 'upstream/main' into dev1

Author: youxiao

.github/workflows/ci.yml

@@ -5,9 +5,25 @@ on:
     branches: [ "main" ]
   pull_request:
     branches: [ "main" ]
+    types: [opened, synchronize, reopened, labeled]
 
 jobs:
+  init-label:
+    if: github.event.action == 'opened'
+    runs-on: ubuntu-latest
+    permissions:
+      pull-requests: write
+    steps:
+      - name: Add run-ci label
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: gh pr edit ${{ github.event.pull_request.number }} --add-label "run-ci"
+
   build:
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     strategy:
       matrix:
@@ -188,7 +204,71 @@ jobs:
         python scripts/test_tensor_api.py -n 1
       shell: bash
 
+  test-sglang-integration:
+    needs: build-flags
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
+    runs-on: ubuntu-latest
+    env:
+      tone_user_name: ${{ secrets.TONE_USER_NAME }}
+    steps:
+      - name: trigger T-one test
+        if: ${{ env.tone_user_name != '' }}
+        run: |
+          curl -L -H "Accept: application/vnd.github+json" -H "X-GitHub-Api-Version: 2022-11-28" https://api.github.com/repos/${{ github.repository }}/actions/runs/${{ github.run_id }}/artifacts > artifact.json
+          cat artifact.json
+          artifact_id=$(jq -r ".artifacts[] | select(.name | contains(\"py312\") ) | .id" artifact.json)
+          signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
+          signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
+          curl -s -H 'Content-Type: application/json' -X POST -d "{\"workspace\":\"mooncake_test\",\"project\":\"mooncake-ci\",\"template\":\"mooncake-ci-test\",\"name\":\"mooncake-ci-${{ github.sha }}\",\"username\":\"${{ secrets.TONE_USER_NAME }}\",\"env_ifs\":\" \",\"env_info\":\"ARTIFACT_ID=${artifact_id} GIT_REPO=${{ github.repository }}\",\"signature\":\"$signature\"}" https://tone.openanolis.cn/api/job/create/ > job.json
+          if [ "$(jq .code job.json)" == 200 ]; then
+            echo "job created"
+          else
+            echo "job create failed"
+            exit 1
+          fi
+          job_id=$(jq .data.id job.json)
+          echo "check job status here and remember to cancel it before restart the job !"
+          echo "job_url: https://tone.openanolis.cn/ws/gclfnh19/test_result/${job_id}"
+          echo "job_id=${job_id}" >> $GITHUB_ENV
+        shell: bash
+  
+      - name: qurey job results
+        if: ${{ env.tone_user_name != '' }}
+        run: |
+          time=0
+          while true; do
+            if [ $time -gt 720 ]; then
+              echo "timeout"
+              exit 1
+            fi
+            signature="${{ secrets.TONE_USER_NAME }}|${{ secrets.TONE_USER_TOKEN }}|$(python3 -c "import time;print(time.time())")"
+            signature="$(python3 -c "import base64;print(base64.b64encode(\"$signature\".encode('utf-8')).decode('utf-8'))")"
+            curl -s -H 'Content-Type: application/json' -X POST -d "{\"username\":\"${{ secrets.TONE_USER_NAME }}\", \"signature\":\"$signature\", \"job_id\": \"${job_id}\"}" https://tone.openanolis.cn/api/job/query/ > job_status.json
+            if ! [ "$(jq .code job_status.json)" == 200 ]; then
+              echo "job query failed"
+              exit 1
+            fi
+            job_status=$(jq .data.job_second_state job_status.json)
+            if [[ $job_status =~ "pass" ]]; then
+              echo "job successful !"
+              exit 0
+            elif [[ $job_status =~ "fail" ]] ; then
+              echo "job failed or stopped !"
+              exit 1
+            fi
+            time=$(( time + 1))
+            sleep 10
+          done
+        shell: bash
+
   build-flags:
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     strategy:
       matrix:
@@ -331,6 +411,10 @@ jobs:
 
   build-docker:
     name: Build Docker Image
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     steps:
     - uses: actions/checkout@v4
@@ -343,6 +427,10 @@ jobs:
 
   spell-check:
     name: Spell Check with Typos
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     steps:
     - name: Checkout Actions Repository
@@ -352,6 +440,10 @@ jobs:
 
   clang-format:
     name: Check code format
+    if: >-
+      github.event_name == 'push' ||
+      github.event.action == 'opened' ||
+      contains(github.event.pull_request.labels.*.name, 'run-ci')
     runs-on: ubuntu-22.04
     steps:
     - name: Checkout Actions Repository

README.md

@@ -243,6 +243,14 @@ The above presents two samples from our trace dataset. The trace includes the ti
 Please kindly cite our paper if you find the paper or the traces are useful:
 
 ```bibtex
+@article{qin2024mooncake,
+  title={Mooncake: A kvcache-centric disaggregated architecture for llm serving},
+  author={Qin, Ruoyu and Li, Zheming and He, Weiran and Cui, Jialei and Tang, Heyi and Ren, Feng and Ma, Teng and Cai, Shangming and Zhang, Yineng and Zhang, Mingxing and others},
+  journal={ACM Transactions on Storage},
+  year={2024},
+  publisher={ACM New York, NY}
+}
+
 @article{qin2024mooncake,
   title        = {Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving},
   author       = {Ruoyu Qin, Zheming Li, Weiran He, Mingxing Zhang, Yongwei Wu, Weimin Zheng, and Xinran Xu},

doc/en/mooncake-store.md

@@ -732,3 +732,13 @@ We provide a reference example `distributed_object_store_provider.py`, located i
 
 #### C++ Usage Example
 The C++ API of Mooncake Store provides more low-level control capabilities. We provide a reference example `client_integration_test`, located in the `mooncake-store/tests` directory. To check if the related components are properly installed, you can run etcd and Master Service (`mooncake_master`) on the same server, and then execute this C++ program (located in the `build/mooncake-store/tests` directory). It should output a successful test result.
+
+## Version Management Policy
+
+The current version of Mooncake Store is defined in [`CMakeLists.txt`](../../mooncake-store/CMakeLists.txt) as `project(MooncakeStore VERSION 2.0.0)`.
+
+When to bump the version:
+
+* **Major version (X.0.0)**: For breaking API changes, major architectural changes, or significant new features that affect backward compatibility
+* **Minor version (0.X.0)**: For new features, API additions, or notable improvements that maintain backward compatibility
+* **Patch version (0.0.X)**: For bug fixes, performance optimizations, or minor improvements that don't affect the API

doc/zh/mooncake-store.md

@@ -738,3 +738,13 @@ python -m mooncake.mooncake_store_service --config=[config_path] --port=8081
 #### C++ 使用示例
 
 Mooncake Store 的 C++ API 提供了更底层的控制能力。我们提供一个参考样例 `client_integration_test`，其位于 `mooncake-store/tests` 目录下。为了检测相关组件是否正常安装，可在相同的服务器上运行 etcd、Master Service（`mooncake_master`），并执行该 C++ 程序（位于 `build/mooncake-store/tests` 目录下），应输出测试成功的结果。
+
+## 版本管理策略
+
+Mooncake Store 的当前版本定义在 [`CMakeLists.txt`](../../mooncake-store/CMakeLists.txt) 中，为 `project(MooncakeStore VERSION 2.0.0)`。
+
+何时需要升级版本：
+
+* **主版本号 (X.0.0)**：当存在破坏性的 API 变更、主要架构更改，或影响向后兼容性的重要新功能时
+* **次版本号 (0.X.0)**：当添加新功能、API 扩展，或保持向后兼容性的显著改进时
+* **修订版本号 (0.0.X)**：当进行错误修复、性能优化，或不改变 API 的细微改进时

docs/source/design/mooncake-store.md

@@ -696,3 +696,13 @@ We provide a reference example `distributed_object_store_provider.py`, located i
 
 #### C++ Usage Example
 The C++ API of Mooncake Store provides more low-level control capabilities. We provide a reference example `client_integration_test`, located in the `mooncake-store/tests` directory. To check if the related components are properly installed, you can run etcd and Master Service (`mooncake_master`) on the same server, and then execute this C++ program (located in the `build/mooncake-store/tests` directory). It should output a successful test result.
+
+## Version Management Policy
+
+The current version of Mooncake Store is defined in [`CMakeLists.txt`](../../mooncake-store/CMakeLists.txt) as `project(MooncakeStore VERSION 2.0.0)`.
+
+When to bump the version:
+
+* **Major version (X.0.0)**: For breaking API changes, major architectural changes, or significant new features that affect backward compatibility
+* **Minor version (0.X.0)**: For new features, API additions, or notable improvements that maintain backward compatibility
+* **Patch version (0.0.X)**: For bug fixes, performance optimizations, or minor improvements that don't affect the API

mooncake-common/common.cmake

@@ -160,14 +160,15 @@ if (USE_ASCEND OR USE_ASCEND_DIRECT)
     file(GLOB ASCEND_TOOLKIT_ROOT "/usr/local/Ascend/ascend-toolkit/latest/*-linux")
   endif()
   set(ASCEND_LIB_DIR "${ASCEND_TOOLKIT_ROOT}/lib64")
-  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
   set(ASCEND_INCLUDE_DIR "${ASCEND_TOOLKIT_ROOT}/include")
   add_compile_options(-Wno-ignored-qualifiers)
   include_directories(/usr/local/include /usr/include ${ASCEND_INCLUDE_DIR})
-  link_directories(${ASCEND_LIB_DIR} ${ASCEND_DEVLIB_DIR})
+  link_directories(${ASCEND_LIB_DIR})
 endif()
 
 if (USE_ASCEND)
+  set(ASCEND_DEVLIB_DIR "${ASCEND_TOOLKIT_ROOT}/devlib")
+  link_directories(${ASCEND_DEVLIB_DIR})
   add_compile_definitions(USE_ASCEND)
 endif()
 

mooncake-ep/include/mooncake_backend.h

@@ -69,6 +69,8 @@ class MooncakeBackend final : public ::c10d::Backend {
         return matrix[location].preferred_hca[0];
     }
 
+    at::Tensor getActiveRanksTensor() { return meta_.activeRanksTensor; }
+
    private:
     static TransferEngine engine_;
     static Transport* transport_;

mooncake-ep/include/mooncake_ep_buffer.h

@@ -81,6 +81,7 @@ struct MooncakeEpBuffer {
     void* qp_devctxs = nullptr;
     std::string device_name;
     bool is_roce_ = false;
+    bool ibgda_disabled_ = false;
 
     // Stream for communication
     at::cuda::CUDAStream comm_stream;
@@ -114,7 +115,9 @@ struct MooncakeEpBuffer {
     torch::Tensor get_next_combine_buffer(int num_max_dispatch_tokens_per_rank,
                                           int hidden, int num_experts);
 
-    void init_ibgda();
+    int init_ibgda();
+
+    bool ibgda_disabled() { return ibgda_disabled_; }
 
     bool is_roce() { return is_roce_; }
 

mooncake-ep/src/mooncake_ep_buffer.cpp

@@ -19,7 +19,13 @@ MooncakeEpBuffer::MooncakeEpBuffer(int rank, int num_ranks,
     CUDA_CHECK(cudaMalloc(&rkeys, num_ranks * sizeof(uint32_t)));
     CUDA_CHECK(
         cudaMalloc(&qp_devctxs, MAX_QP_COUNT * sizeof(mlx5gda_qp_devctx)));
-    init_ibgda();
+    int ret = init_ibgda();
+    if (ret != 0) {
+        LOG(WARNING) << "Failed to initialize IBGDA. "
+                     << "Using fallback implementation. "
+                     << "Performance will be degraded.";
+        ibgda_disabled_ = true;
+    }
 
     // Create 32 MiB workspace
     CUDA_CHECK(cudaMalloc(&workspace, NUM_WORKSPACE_BYTES));
@@ -279,7 +285,7 @@ torch::Tensor MooncakeEpBuffer::get_next_combine_buffer(
         torch::TensorOptions().dtype(dtype).device(torch::kCUDA));
 }
 
-void MooncakeEpBuffer::init_ibgda() {
+int MooncakeEpBuffer::init_ibgda() {
     int num_devices;
     ibv_device** dev_list = ibv_get_device_list(&num_devices);
     int nic_id = -1;
@@ -299,7 +305,7 @@ void MooncakeEpBuffer::init_ibgda() {
     ibv_context* ctx = ibv_open_device(dev_list[nic_id]);
     if (!ctx) {
         perror("Failed to open device");
-        exit(1);
+        return -1;
     }
     if (ibv_query_gid(ctx, 1, 3, &gid)) {
         perror("Failed to query gid");
@@ -309,7 +315,7 @@ void MooncakeEpBuffer::init_ibgda() {
     ibv_pd* pd = ibv_alloc_pd(ctx);
     if (!pd) {
         perror("Failed to allocate protection domain");
-        exit(1);
+        return -1;
     }
     mlx5dv_pd mpd;
     mlx5dv_obj dv_obj = {};
@@ -334,24 +340,24 @@ void MooncakeEpBuffer::init_ibgda() {
         fprintf(stderr,
                 "If the error is `Bad address`, probably because your GPU "
                 "does not support GPUDirect RDMA.\n");
-        exit(1);
+        return -1;
     }
     memheap* ctrl_buf_heap = memheap_create(CTRL_BUF_SIZE);
     if (!ctrl_buf_heap) {
         perror("Failed to create memory heap");
-        exit(1);
+        return -1;
     }
     for (int i = 0; i < MAX_QP_COUNT; ++i) {
         mlx5gda_qp* qp = mlx5gda_create_rc_qp(mpd, ctrl_buf, ctrl_buf_umem,
                                               ctrl_buf_heap, pd, 16384, 1);
         if (!qp) {
             perror("Failed to create QP");
-            exit(1);
+            return -1;
         }
         is_roce_ = qp->port_attr.link_layer == IBV_LINK_LAYER_ETHERNET;
         if (mlx5gda_modify_rc_qp_rst2init(qp, 0)) {
             perror("Failed to mlx5gda_modify_rc_qp_rst2init");
-            exit(1);
+            return -1;
         }
         mlx5gda_qp_devctx qp_devctx = {
             .qpn = qp->qpn,
@@ -365,6 +371,7 @@ void MooncakeEpBuffer::init_ibgda() {
                    sizeof(mlx5gda_qp_devctx), cudaMemcpyHostToDevice);
         qps.push_back(qp);
     }
+    return 0;
 }
 
 void MooncakeEpBuffer::sync_ib(const std::vector<int64_t>& remote_addrs,

mooncake-integration/ep/ep_py.cpp

@@ -57,12 +57,19 @@ std::string getPreferredHca(c10::intrusive_ptr<c10d::Backend> backend,
     return mooncakeBackend->getPreferredHca(location);
 }
 
+at::Tensor getActiveRanks(c10::intrusive_ptr<c10d::Backend> backend) {
+    auto mooncakeBackend =
+        c10::static_intrusive_pointer_cast<MooncakeBackend>(backend);
+    return mooncakeBackend->getActiveRanksTensor();
+}
+
 PYBIND11_MODULE(ep, m) {
     m.def("createMooncakeBackend", &createMooncakeBackend);
     m.def("createMooncakeCpuBackend", &createMooncakeCpuBackend);
     m.def("set_host_ip", &MooncakeBackend::setHostIp);
     m.def("set_device_filter", &MooncakeBackend::setDeviceFilter);
     m.def("get_preferred_hca", &getPreferredHca);
+    m.def("get_active_ranks", &getActiveRanks);
 
     py::class_<MooncakeBackend::MooncakeBackendOptions,
                c10::intrusive_ptr<MooncakeBackend::MooncakeBackendOptions>>(
@@ -79,6 +86,7 @@ PYBIND11_MODULE(ep, m) {
 
     py::class_<MooncakeEpBuffer>(m, "Buffer")
         .def(py::init<int, int, int64_t, std::string>())
+        .def("ibgda_disabled", &MooncakeEpBuffer::ibgda_disabled)
         .def("is_roce", &MooncakeEpBuffer::is_roce)
         .def("sync_ib", &MooncakeEpBuffer::sync_ib)
         .def("sync_roce", &MooncakeEpBuffer::sync_roce)