# EXPERIMENT: arx RDMA all-reduce for decode-sized TP all-reduces (see arx.py),
# and with VLLM_ARXBIG=1 arxbig.cu's all-gather for prefill-sized SP gathers.
#   docker compose -f compose/glm53.yaml -f compose/arx.yaml up -d
services:
  glm53:
    volumes:
      - ../experimental/arx/arx.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx.py:ro
      - ../experimental/arx/arx_vllm.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arx_vllm.cu:ro
      - ../experimental/arx/arxbig.cu:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/arxbig.cu:ro
      - ../experimental/arx/cuda_communicator.py:/usr/local/lib/python3.12/dist-packages/vllm/distributed/device_communicators/cuda_communicator.py:ro
    environment:
      - VLLM_ARX_ALLREDUCE=1
      - VLLM_ARXBIG=${VLLM_ARXBIG:-1}
      - VLLM_ARXBIG_RS=${VLLM_ARXBIG_RS:-1}
      - VLLM_ARXBIG_AG=${VLLM_ARXBIG_AG:-0}
      # Decode all-reduces above 256 KB and up to this size run arx's two-shot mode
      # (reduce-scatter + all-gather over RDMA); 0 sends them to NCCL.
      - VLLM_ARX_TWO_SHOT_MAX_KB=${VLLM_ARX_TWO_SHOT_MAX_KB:-2048}