Skip to content

[TransferEngine][Store] Fix and complete Biren (SUPA) build and device integration - #4317

Open
Dashener2 wants to merge 2 commits into
kvcache-ai:mainfrom
Dashener2:feat/biren-supa-integration
Open

Dashener2 wants to merge 2 commits into
kvcache-ai:mainfrom
Dashener2:feat/biren-supa-integration

Conversation

@Dashener2

Copy link
Copy Markdown

SUPA support was only wired into rdma_context and cuda_alike; every other vendor-conditional omitted USE_SUPA, and gpu_vendor/supa.h mapped CUDA names to non-existent supa* symbols, so it never compiled against a real SUPA SDK.

Transfer Engine:

  • Discover GPU topology and classify device memory for SUPA (topology.cpp, memory_location.cpp).
  • Enable the TCP GPU staging path for SUPA.
  • Build examples and the TCP transport test with SUPA.
  • Rewrite gpu_vendor/supa.h against the real SUPA API: runtime symbols use the su prefix, driver symbols use sudrv, driver enums use their su* names, and suMemcpyAsync takes the stream before the kind (wrapped to preserve the CUDA argument order).

Mooncake Store:

  • Add AcceleratorVendor::kSupa and register the CUDA-like accelerator device.

Build:

  • Include the SUPA SDK as a system include directory so its generic headers (e.g. version.h) do not shadow Mooncake headers.
  • Discover the runtime library by name (libsupa/libsupa-runtime) instead of hardcoding it, and link the resolved libraries.

Docs: document -DUSE_SUPA and BIREN_HOME in build.md, quick-start.md and supported-protocols.md.

Tested on 8x Biren166M + 8x RoCE: full TE and Store build and link with -DUSE_SUPA=ON against the SUPA SDK (GPU allocation verified at runtime); the default build without USE_SUPA is unchanged.

Description

Biren (SUPA) support existed only as a skeleton: rdma_context.cpp and
cuda_alike.h knew about USE_SUPA, but every other vendor-conditional omitted
it, and gpu_vendor/supa.h mapped the CUDA API onto supa* symbols that do not
exist in the real SUPA SDK (suGetDeviceCount, suMallocDevice, sudrvMemCreate,
...). As a result the -DUSE_SUPA=ON build could not compile, and even if it had,
GPU topology discovery, device-memory classification, the TCP staging path, and
the Store accelerator device would all have fallen back to CPU paths.

This change completes the Biren integration so -DUSE_SUPA=ON builds and works
end to end:

  • Transfer Engine
    • topology.cpp: include USE_SUPA in the GPU topology guard so supa:N
      entries are discovered and NIC/GPU affinity applies.
    • memory_location.cpp: classify Biren device memory via cuPointerGetAttributes
      instead of treating it as host memory.
    • tcp_transport.cpp / tcp_transport_session_impl.h: enable GPU-aware
      staging for SUPA.
    • example/* and tests/tcp_transport_test.cpp: build the GPU paths under SUPA.
    • gpu_vendor/supa.h: rewrite against the actual SUPA SDK. Runtime API uses
      the su prefix, driver API uses sudrv, driver enums use their su* names,
      NVIDIA-only fabric/VMM concepts are dropped, and suMemcpyAsync (which takes
      the stream before the kind) is wrapped to keep the CUDA call order.
  • Mooncake Store
    • Add AcceleratorVendor::kSupa, register the CUDA-like accelerator device, and
      enable the cuda-like registry path for SUPA.
  • Build
    • Expose ${BIREN_HOME}/supa/include as a system include dir: the SDK ships
      generic headers (notably version.h) that otherwise shadow Mooncake's own
      generated version.h.
    • Resolve the runtime library by name (libsupa for older SDKs,
      libsupa-runtime for current ones) and link the resolved library set.
  • Docs: add Biren rows to build.md, quick-start.md and
    supported-protocols.md, documenting -DUSE_SUPA=ON and BIREN_HOME.

No related issue; the work aligns with the "Domestic GPU Support" milestone of
the overall roadmap.

Module

  • Transfer Engine (mooncake-transfer-engine)
  • Mooncake Store (mooncake-store)
  • Mooncake Conductor (mooncake-conductor)
  • Reshard (mooncake-reshard)
  • Mooncake EP (mooncake-ep)
  • Mooncake PG (mooncake-pg)
  • Integration (mooncake-integration)
  • P2P Store (mooncake-p2p-store)
  • Python Wheel (mooncake-wheel)
  • Common (mooncake-common)
  • Mooncake RL (mooncake-rl)
  • CI/CD
  • Docs
  • Other

Type of Change

  • Bug fix
  • New feature
  • Refactor
  • Breaking change
  • Documentation update
  • Performance improvement
  • Other

How Has This Been Tested?

Tested inside a Biren SUPA dev container on a host with 8x Biren166M (64 GB each),
Hygon C86-4G CPU and 8x RoCE NICs, using SUPA SDK 1.10.

Test commands:

# 1. SUPA build + link (TE, Store, examples, tests)
cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \
      -DUSE_SUPA=ON -DBIREN_HOME=/usr/local/birensupa/sdk/latest \
      -DWITH_STORE_RUST=OFF
ninja -C build transfer_engine mooncake_store \
      transfer_engine_bench transfer_engine_validator \
      transfer_engine_bench_with_notify tcp_transport_test rdma_transport_test

# 2. Default build regression (no SUPA)
cmake -S . -B build_default -G Ninja -DCMAKE_BUILD_TYPE=Release -DWITH_STORE_RUST=OFF
ninja -C build_default transfer_engine mooncake_store

# 3. Runtime smoke test: allocate GPU memory on a Biren166M via the TCP path
./build/mooncake-transfer-engine/example/transfer_engine_bench \
  --mode=target --metadata_server=P2PHANDSHAKE --local_server_name=127.0.0.1:12345 \
  --protocol=tcp --use_vram=true --gpu_id=0 \
  --buffer_size=$((256*1024*1024)) --duration=8

Test results:

  • Unit tests pass

  • Integration tests pass (if applicable)

  • Manual testing done (describe below)

  • -DUSE_SUPA=ON: 0 compiler errors. libtransfer_engine.a and
    libmooncake_store.a link, and all examples/tests link; the binaries resolve
    libsupa-runtime.so.1 at runtime.

  • Runtime smoke test: the target allocates a 256 MiB buffer in Biren VRAM
    ("VRAM is used" / "Allocating memory on GPU 0") and the TCP transport starts
    cleanly, confirming the su* API mapping and the new memory_location /
    topology paths compile and run.

  • Default build (USE_SUPA=OFF): 0 compiler errors, no regression to
    existing backends.

  • Not covered here: RDMA end-to-end transfer could not be exercised inside the
    container because the zrdma/dinghai10e NIC userspace provider is not
    installed there (ibv_devices is empty), and rdma_transport_test /
    tcp_transport_test were built but not executed. A follow-up run on a bare
    host (or an image with the vendor RDMA provider) is recommended before
    merging.

Checklist

  • I have performed a self-review of my own code
  • I have formatted my code using ./scripts/code_format.sh
  • I have run pre-commit on the files changed in this PR and all hooks pass
  • I have updated the documentation (if applicable)
  • I have added tests to prove my changes are effective
  • For changes >500 LOC: I have filed an RFC issue

Notes: pre-commit run --files ... passes on all changed files
(trailing-whitespace, end-of-file-fixer, check-merge-conflict,
check-added-large-files, clang-format, codespell and cmake-format); the Python
and Rust hooks are skipped because no Python/Rust files are touched. No new unit
tests were added; the change is compile-gated enablement, and the existing
tcp_transport_test / rdma_transport_test / examples now build under
USE_SUPA. The change is ~264 insertions, under the 500 LOC RFC threshold.

AI Assistance Disclosure

  • No AI tools were used
  • AI tools were used (specify below)

An AI coding assistant (opencode, deepseek-flash) was used to locate the missing
USE_SUPA branches, rewrite gpu_vendor/supa.h against the real SUPA SDK
symbols, adjust the CMake, and run the containerized build/regression
validation. All changed lines were reviewed by the submitter, who can defend the
change end to end.

SUPA support was only wired into rdma_context and cuda_alike; every other
vendor-conditional omitted USE_SUPA, and gpu_vendor/supa.h mapped CUDA names
to non-existent `supa*` symbols, so it never compiled against a real SUPA SDK.

Transfer Engine:
- Discover GPU topology and classify device memory for SUPA (topology.cpp,
  memory_location.cpp).
- Enable the TCP GPU staging path for SUPA.
- Build examples and the TCP transport test with SUPA.
- Rewrite gpu_vendor/supa.h against the real SUPA API: runtime symbols use the
  `su` prefix, driver symbols use `sudrv`, driver enums use their `su*` names,
  and suMemcpyAsync takes the stream before the kind (wrapped to preserve the
  CUDA argument order).

Mooncake Store:
- Add AcceleratorVendor::kSupa and register the CUDA-like accelerator device.

Build:
- Include the SUPA SDK as a system include directory so its generic headers
  (e.g. version.h) do not shadow Mooncake headers.
- Discover the runtime library by name (libsupa/libsupa-runtime) instead of
  hardcoding it, and link the resolved libraries.

Docs: document -DUSE_SUPA and BIREN_HOME in build.md, quick-start.md and
supported-protocols.md.

Tested on 8x Biren166M + 8x RoCE: full TE and Store build and link with
-DUSE_SUPA=ON against the SUPA SDK (GPU allocation verified at runtime); the
default build without USE_SUPA is unchanged.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Common documentation Improvements or additions to documentation run-ci Store Transfer Engine

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant