diff --git a/.github/workflows/conda-ci.yml b/.github/workflows/conda-ci.yml index 8fdedb8f7f..6cb0f5ce2b 100644 --- a/.github/workflows/conda-ci.yml +++ b/.github/workflows/conda-ci.yml @@ -14,7 +14,7 @@ jobs: runs-on: self-hosted container: image: lmsysorg/sglang:v0.5.0rc0-cu126 - options: --gpus all --ipc=host --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --memory=0 --memory-swap=0 -v /mnt/nvme0n1/models:/root/models -v /mnt/nvme0n1/datasets:/root/datasets + options: --privileged --cap-add SYS_NICE --security-opt seccomp=unconfined --gpus all --ipc=host --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --memory=0 --memory-swap=0 -v /mnt/nvme0n1/models:/root/models -v /mnt/nvme0n1/datasets:/root/datasets defaults: run: diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 9b485a5ade..f88e7cf451 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -31,28 +31,11 @@ jobs: runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}] + info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_ppo_critic_only_short.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -69,23 +52,61 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-sglang-config: @@ -93,23 +114,6 @@ jobs: runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false @@ -131,23 +135,61 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-megatron: @@ -155,28 +197,11 @@ jobs: runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: - info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}] + info: [{"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -193,23 +218,61 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-precision: @@ -217,23 +280,6 @@ jobs: runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false @@ -255,23 +301,61 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-ckpt: @@ -279,28 +363,11 @@ jobs: runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: - info: [{"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py --async-save"}] + info: [{"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -317,23 +384,61 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-plugin-contracts: @@ -382,44 +487,35 @@ jobs: - name: Execute shell: bash run: | + TEST_PATH="${{ matrix.info.test_file }}" if [[ "$TEST_PATH" != tests/* ]]; then TEST_PATH="tests/$TEST_PATH" fi + TEST_ARGS="${{ matrix.info.test_args || '' }}" + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" + python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" fi + e2e-test-image: if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-image')) runs-on: self-hosted - container: - image: slimerl/slime-test:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py --async-save"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}] + info: [{"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py"}, {"num_gpus": 8, "test_file": "test_qwen3_30B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -436,38 +532,66 @@ jobs: uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime-test:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' + e2e-test-changed-detect: if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-changed')) runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 outputs: matrix: ${{ steps.detect.outputs.matrix }} has_tests: ${{ steps.detect.outputs.has_tests }} @@ -508,23 +632,6 @@ jobs: needs: e2e-test-changed-detect if: needs.e2e-test-changed-detect.outputs.has_tests == 'true' runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: ${{ fromJson(needs.e2e-test-changed-detect.outputs.matrix) }} @@ -543,19 +650,55 @@ jobs: - name: Checkout repository uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi \ No newline at end of file + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' \ No newline at end of file diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 4f88e5d02c..16cc38bd32 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -2,8 +2,8 @@ 'e2e-test-short': { 'label': 'run-ci-short', 'tests': [ - {'test_file': 'test_qwen2.5_0.5B_gsm8k_async_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen2.5_0.5B_gsm8k_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4}, {'test_file': 'test_qwen2.5_0.5B_ppo_critic_only_short.py', 'num_gpus': 4}, ], }, @@ -20,7 +20,9 @@ 'label': 'run-ci-megatron', 'tests': [ {'test_file': 'test_quick_start_glm4_9B.py', 'num_gpus': 8}, + {'test_file': 'test_glm4.7_30B_A3B_pd_mooncake.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1'}, + {'test_file': 'test_qwen3_30B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'}, {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'use_deepep': '1', 'use_fp8_rollout': '1', 'enable_eval': '0'}, {'test_file': 'test_qwen3_30B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'}, {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, @@ -42,7 +44,7 @@ 'label': 'run-ci-ckpt', 'tests': [ {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ckpt.py --async-save', 'num_gpus': 8}, + {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8}, ], }, @@ -62,16 +64,18 @@ 'label': 'run-ci-image', 'image': 'slimerl/slime-test:latest', 'tests': [ - {'test_file': 'test_qwen2.5_0.5B_gsm8k_async_short.py', 'num_gpus': 4}, - {'test_file': 'test_qwen2.5_0.5B_gsm8k_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen3.5_0.8B_gsm8k_async_short.py', 'num_gpus': 4}, + {'test_file': 'test_qwen3.5_0.8B_gsm8k_short.py', 'num_gpus': 4}, {'test_file': 'test_quick_start_glm4_9B.py', 'num_gpus': 8}, + {'test_file': 'test_glm4.7_30B_A3B_pd_mooncake.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_30B_A3B.py', 'num_gpus': 8}, + {'test_file': 'test_qwen3_30B_A3B_pd_mooncake.py', 'num_gpus': 8, 'use_deepep': '1'}, {'test_file': 'test_qwen3_4B_ppo.py', 'num_gpus': 8}, {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8}, {'test_file': 'test_mimo_7B_mtp_only_grad.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8}, {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8}, - {'test_file': 'test_qwen3_4B_ckpt.py --async-save', 'num_gpus': 8}, + {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8}, ], @@ -110,23 +114,6 @@ jobs: runs-on: ubuntu-latest <% else %> runs-on: self-hosted - container: - image: << config.image if config.image else 'slimerl/slime:latest' >> - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets <% endif %> strategy: fail-fast: false @@ -164,39 +151,83 @@ jobs: shell: bash run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps <% else %> - - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages <% endif %> - name: Execute shell: bash run: | +<% if config.get('cpu') %> TEST_PATH="${{ matrix.info.test_file }}" if [[ "$TEST_PATH" != tests/* ]]; then TEST_PATH="tests/$TEST_PATH" fi + TEST_ARGS="${{ matrix.info.test_args || '' }}" + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf '%s\n' "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" + python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" fi +<% else %> + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + << config.image if config.image else 'slimerl/slime:latest' >> \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' +<% endif %> <% endfor %> e2e-test-changed-detect: if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-changed')) runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 outputs: matrix: ${{ steps.detect.outputs.matrix }} has_tests: ${{ steps.detect.outputs.has_tests }} @@ -237,23 +268,6 @@ jobs: needs: e2e-test-changed-detect if: needs.e2e-test-changed-detect.outputs.has_tests == 'true' runs-on: self-hosted - container: - image: slimerl/slime:latest - options: > - --gpus all - --ipc=host - --shm-size=16g - --ulimit memlock=-1 - --ulimit stack=67108864 - --memory=0 - --memory-swap=0 - -e http_proxy=$http_proxy - -e https_proxy=$https_proxy - -e HTTP_PROXY=$HTTP_PROXY - -e HTTPS_PROXY=$HTTPS_PROXY - -v /mnt/nvme0n1/slime_ci:/data/slime_ci - -v /mnt/nvme0n1/slime_ci/models:/root/models - -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets strategy: fail-fast: false matrix: ${{ fromJson(needs.e2e-test-changed-detect.outputs.matrix) }} @@ -272,19 +286,55 @@ jobs: - name: Checkout repository uses: actions/checkout@v4 - - name: Install - shell: bash - run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages - - name: Execute shell: bash run: | - TEST_PATH="${{ matrix.info.test_file }}" - if [[ "$TEST_PATH" != tests/* ]]; then - TEST_PATH="tests/$TEST_PATH" - fi - if [ "${{ matrix.info.num_gpus }}" = "0" ]; then - python "$TEST_PATH" - else - python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" - fi + docker run --rm \ + --privileged \ + --cap-add SYS_NICE \ + --security-opt seccomp=unconfined \ + --network host \ + --gpus all \ + --ipc=host \ + --shm-size=16g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --memory=0 \ + --memory-swap=0 \ + -e http_proxy \ + -e https_proxy \ + -e HTTP_PROXY \ + -e HTTPS_PROXY \ + -e GITHUB_COMMIT_NAME \ + -e WANDB_API_KEY \ + -e SLIME_TEST_ENABLE_INFINITE_RUN \ + -e SLIME_TEST_USE_DEEPEP \ + -e SLIME_TEST_USE_FP8_ROLLOUT \ + -e SLIME_TEST_ENABLE_EVAL \ + -e TEST_FILE="${{ matrix.info.test_file }}" \ + -e TEST_ARGS="${{ matrix.info.test_args || '' }}" \ + -e NUM_GPUS="${{ matrix.info.num_gpus }}" \ + -v "$GITHUB_WORKSPACE:$GITHUB_WORKSPACE" \ + -v /mnt/nvme0n1/slime_ci:/data/slime_ci \ + -v /mnt/nvme0n1/slime_ci/models:/root/models \ + -v /mnt/nvme0n1/slime_ci/datasets:/root/datasets \ + -w "$GITHUB_WORKSPACE" \ + slimerl/slime:latest \ + bash -lc ' + set -euo pipefail + pip install -e . --no-deps --break-system-packages + TEST_PATH="$TEST_FILE" + if [[ "$TEST_PATH" != tests/* ]]; then + TEST_PATH="tests/$TEST_PATH" + fi + if [[ -n "$TEST_ARGS" ]]; then + read -r -a TEST_ARGS_ARRAY < <(printf "%s\n" "$TEST_ARGS") + else + TEST_ARGS_ARRAY=() + fi + if [ "$NUM_GPUS" = "0" ]; then + python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + else + python tests/ci/gpu_lock_exec.py --count "$NUM_GPUS" -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}" + fi + ' diff --git a/docker/Dockerfile b/docker/Dockerfile index f19befd7b5..a4791f2548 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,4 +1,4 @@ -ARG SGLANG_IMAGE_TAG=v0.5.9 +ARG SGLANG_IMAGE_TAG=v0.5.10.post1 FROM slimerl/sglang:${SGLANG_IMAGE_TAG} AS sglang # ======================================== Arguments ============================================= @@ -64,7 +64,8 @@ RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ fi COPY requirements.txt /tmp/requirements.txt -RUN pip install -r /tmp/requirements.txt +RUN pip install --ignore-installed PyJWT && \ + pip install -r /tmp/requirements.txt # Temporarily install another sgl-kernel version for GB300 without rebuilding the whole image RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \ diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch index e9145702e1..4a13e2f9b4 100644 --- a/docker/patch/latest/sglang.patch +++ b/docker/patch/latest/sglang.patch @@ -1,8 +1,8 @@ diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py -index 6fbd1db823..f80ec11bb4 100644 +index 691f06411d..671ac81c48 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py -@@ -274,6 +274,7 @@ class ModelConfig: +@@ -294,6 +294,7 @@ class ModelConfig: if is_draft_model and self.hf_config.architectures[0] in [ "DeepseekV3ForCausalLM", @@ -11,7 +11,7 @@ index 6fbd1db823..f80ec11bb4 100644 ]: self.hf_config.architectures[0] = "DeepseekV3ForCausalLMNextN" diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py -index da4629e525..c03f98231a 100644 +index f7d4092d85..3aae51c849 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -17,6 +17,7 @@ class KVArgs: @@ -22,45 +22,8 @@ index da4629e525..c03f98231a 100644 aux_data_ptrs: List[int] aux_data_lens: List[int] aux_item_lens: List[int] -diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py -index 67fe82ad67..ed5fa7b0e3 100644 ---- a/python/sglang/srt/disaggregation/common/conn.py -+++ b/python/sglang/srt/disaggregation/common/conn.py -@@ -333,6 +333,10 @@ class CommonKVReceiver(BaseKVReceiver): - self.required_dst_info_num = ( - self.kv_mgr.attn_tp_size // self.prefill_attn_tp_size - ) -+ # With attention DP, one request is routed to one decode rank. -+ # Waiting for all TP shards to pre-allocate the same bootstrap room would stall forever. -+ if self.kv_mgr.attn_dp_size > 1: -+ self.required_dst_info_num = 1 - self.required_prefill_response_num = 1 * ( - self.prefill_pp_size // self.kv_mgr.pp_size - ) -@@ -422,6 +426,7 @@ class CommonKVReceiver(BaseKVReceiver): - f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank} and target_dp_group: {self.target_dp_group} and target_pp_rank {target_pp_rank}", - ) - self.kv_mgr.update_status(self.bootstrap_room, KVPoll.Failed) -+ self.bootstrap_infos = None - return - - self.bootstrap_infos = bootstrap_infos -@@ -610,8 +615,12 @@ class CommonKVBootstrapServer(BaseKVBootstrapServer): - and int(target_dp_group) == -1 - and int(target_pp_rank) == -1 - ): -+ inferred_attn_tp_size = max( -+ (len(v) for v in self.prefill_port_table.values()), -+ default=self.attn_tp_size, -+ ) - prefill_parallel_info = { -- "prefill_attn_tp_size": self.attn_tp_size, -+ "prefill_attn_tp_size": inferred_attn_tp_size, - "prefill_dp_size": self.dp_size, - "prefill_pp_size": self.pp_size, - "prefill_page_size": self.page_size, diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py -index 1d8baf0028..1ebb959298 100644 +index f54c882cc2..03832002f0 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -21,6 +21,7 @@ Life cycle of a request in the decode server @@ -71,7 +34,7 @@ index 1d8baf0028..1ebb959298 100644 import time from collections import deque from dataclasses import dataclass -@@ -40,8 +41,10 @@ from sglang.srt.disaggregation.utils import ( +@@ -42,8 +43,10 @@ from sglang.srt.disaggregation.utils import ( MetadataBuffers, ReqToMetadataIdxAllocator, TransferBackend, @@ -81,8 +44,8 @@ index 1d8baf0028..1ebb959298 100644 + is_slime_profiling_enabled, kv_to_page_indices, poll_and_all_reduce, - prepare_abort, -@@ -295,6 +298,7 @@ class DecodePreallocQueue: + poll_and_all_reduce_with_staging, +@@ -344,6 +347,7 @@ class DecodePreallocQueue: kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( self.metadata_buffers.get_buf_infos() ) @@ -90,8 +53,8 @@ index 1d8baf0028..1ebb959298 100644 if hasattr(self.token_to_kv_pool, "get_state_buf_infos"): state_data_ptrs, state_data_lens, state_item_lens = ( -@@ -336,6 +340,16 @@ class DecodePreallocQueue: - ) +@@ -398,6 +402,16 @@ class DecodePreallocQueue: + ) return kv_manager + def release_memory_occupation(self): @@ -107,7 +70,7 @@ index 1d8baf0028..1ebb959298 100644 def add(self, req: Req, is_retracted: bool = False) -> None: """Add a request to the pending queue.""" if self._check_if_req_exceed_kv_capacity(req): -@@ -440,12 +454,37 @@ class DecodePreallocQueue: +@@ -525,12 +539,37 @@ class DecodePreallocQueue: [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group ) @@ -145,16 +108,16 @@ index 1d8baf0028..1ebb959298 100644 + self.scheduler.metrics_collector.increment_bootstrap_failed_reqs() elif poll == KVPoll.WaitingForInput: decode_req.waiting_for_input = True - elif poll == KVPoll.Failed: -@@ -590,6 +629,7 @@ class DecodePreallocQueue: + decode_req.req.time_stats.set_bootstrap_done_time() +@@ -770,6 +809,7 @@ class DecodePreallocQueue: self.req_to_metadata_buffer_idx_allocator.alloc() ) assert decode_req.metadata_buffer_index is not None + self.metadata_buffers.clear_profiling_buf(decode_req.metadata_buffer_index) page_indices = kv_to_page_indices(kv_indices, page_size) - decode_req.kv_receiver.init( + decode_req.kv_receiver.send_metadata( page_indices, decode_req.metadata_buffer_index, state_indices -@@ -751,6 +791,7 @@ class DecodeTransferQueue: +@@ -964,6 +1004,7 @@ class DecodeTransferQueue: output_topk_index, output_hidden_states, output_bootstrap_room, @@ -162,7 +125,7 @@ index 1d8baf0028..1ebb959298 100644 ) = self.metadata_buffers.get_buf(idx) # Validate bootstrap_room to detect context corruption -@@ -813,6 +854,14 @@ class DecodeTransferQueue: +@@ -1025,6 +1066,14 @@ class DecodeTransferQueue: output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist() ) @@ -176,10 +139,10 @@ index 1d8baf0028..1ebb959298 100644 + decode_req.kv_receiver.clear() decode_req.kv_receiver = None - trace_slice_end( -@@ -830,6 +879,13 @@ class DecodeTransferQueue: - [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group - ) + decode_req.req.time_stats.set_wait_queue_entry_time() +@@ -1057,6 +1106,13 @@ class DecodeTransferQueue: + [dr.kv_receiver for dr in self.queue], self.gloo_group + ) + # Transfer timeout: if a request has been in the transfer queue for too long + # (e.g., stuck in Bootstrapping/WaitingForInput/Transferring), treat it as failed. @@ -191,7 +154,7 @@ index 1d8baf0028..1ebb959298 100644 transferred_reqs = [] indices_to_remove = set() for i, (decode_req, poll) in enumerate(zip(self.queue, polls)): -@@ -877,7 +933,20 @@ class DecodeTransferQueue: +@@ -1111,7 +1167,20 @@ class DecodeTransferQueue: KVPoll.WaitingForInput, KVPoll.Transferring, ]: @@ -213,7 +176,7 @@ index 1d8baf0028..1ebb959298 100644 else: raise ValueError(f"Unexpected poll case: {poll}") -@@ -893,6 +962,14 @@ class DecodeTransferQueue: +@@ -1132,6 +1201,14 @@ class DecodeTransferQueue: return transferred_reqs @@ -228,7 +191,7 @@ index 1d8baf0028..1ebb959298 100644 class SchedulerDisaggregationDecodeMixin: -@@ -1072,7 +1149,15 @@ class SchedulerDisaggregationDecodeMixin: +@@ -1301,7 +1378,15 @@ class SchedulerDisaggregationDecodeMixin: resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs() self.waiting_queue.extend(resumed_reqs) if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: @@ -246,39 +209,35 @@ index 1d8baf0028..1ebb959298 100644 if not hasattr(self, "polling_count"): diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py -index d0d4efd958..b3a207063e 100644 +index 64d97f5c69..4ef08446aa 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py -@@ -30,7 +30,7 @@ from sglang.srt.disaggregation.common.utils import ( - from sglang.srt.disaggregation.mooncake.utils import ( - check_mooncake_custom_mem_pool_enabled, +@@ -31,6 +31,7 @@ from sglang.srt.disaggregation.mooncake.utils import ( + from sglang.srt.disaggregation.utils import ( + DisaggregationMode, + filter_kv_indices_for_cp_rank, ++ iter_aux_transfer_specs, ) --from sglang.srt.disaggregation.utils import DisaggregationMode -+from sglang.srt.disaggregation.utils import DisaggregationMode, iter_aux_transfer_specs from sglang.srt.distributed.parallel_state import get_mooncake_transfer_engine from sglang.srt.environ import envs - from sglang.srt.server_args import ServerArgs -@@ -260,6 +260,19 @@ class MooncakeKVManager(CommonKVManager): +@@ -276,6 +277,16 @@ class MooncakeKVManager(CommonKVManager): self.kv_args.state_data_ptrs, self.kv_args.state_data_lens ) + def deregister_buffer_to_engine(self): -+ # Batch deregister KV data buffers + if self.kv_args.kv_data_ptrs: + self.engine.batch_deregister(self.kv_args.kv_data_ptrs) + -+ # Batch deregister auxiliary data buffers + if self.kv_args.aux_data_ptrs: + self.engine.batch_deregister(self.kv_args.aux_data_ptrs) + -+ # Batch deregister state/extra pool data buffers + if self.kv_args.state_data_ptrs: + self.engine.batch_deregister(self.kv_args.state_data_ptrs) + - def _transfer_data(self, mooncake_session_id, transfer_blocks): - if not transfer_blocks: - return 0 -@@ -524,10 +537,14 @@ class MooncakeKVManager(CommonKVManager): + # ------------------------------------------------------------------ + # Staging buffer methods (all delegate to staging_handler.py) + # ------------------------------------------------------------------ +@@ -884,10 +895,14 @@ class MooncakeKVManager(CommonKVManager): prefill_aux_ptrs = self.kv_args.aux_data_ptrs prefill_aux_item_lens = self.kv_args.aux_item_lens @@ -297,7 +256,7 @@ index d0d4efd958..b3a207063e 100644 transfer_blocks.append((src_addr, dst_addr, length)) return self._transfer_data(req.mooncake_session_id, transfer_blocks) -@@ -541,9 +558,14 @@ class MooncakeKVManager(CommonKVManager): +@@ -901,9 +916,14 @@ class MooncakeKVManager(CommonKVManager): prefill_aux_ptrs = self.kv_args.aux_data_ptrs prefill_aux_item_lens = self.kv_args.aux_item_lens @@ -315,7 +274,7 @@ index d0d4efd958..b3a207063e 100644 data = AuxDataCodec.serialize_data_from_buffer(src_addr, length) self.send_aux_data_to_endpoint( -@@ -643,13 +665,13 @@ class MooncakeKVManager(CommonKVManager): +@@ -1002,13 +1022,13 @@ class MooncakeKVManager(CommonKVManager): raise RuntimeError( f"PD Disaggregation does NOT support PD different TP sizes for non-MLA {state_type.upper()} hybrid models yet." ) @@ -335,7 +294,7 @@ index d0d4efd958..b3a207063e 100644 # Reuse _send_kvcache_generic interface to send extra pool data prefill_state_indices = np.array(prefill_state_indices, dtype=np.int32) dst_state_indices = np.array(req.dst_state_indices, dtype=np.int32) -@@ -858,12 +880,6 @@ class MooncakeKVManager(CommonKVManager): +@@ -1266,12 +1286,6 @@ class MooncakeKVManager(CommonKVManager): if ret != 0: with self.session_lock: self.session_failures[req.mooncake_session_id] += 1 @@ -347,10 +306,10 @@ index d0d4efd958..b3a207063e 100644 - ) self.record_failure( kv_chunk.room, - f"Failed to send kv chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", -@@ -880,13 +896,31 @@ class MooncakeKVManager(CommonKVManager): + f"Failed to send kv chunk of {kv_chunk.room} to " +@@ -1289,13 +1303,31 @@ class MooncakeKVManager(CommonKVManager): - if kv_chunk.is_last: + if kv_chunk.is_last_chunk: if kv_chunk.state_indices is not None: - self.maybe_send_extra( + ret = self.maybe_send_extra( @@ -361,13 +320,13 @@ index d0d4efd958..b3a207063e 100644 target_rank_registration_info, ) + if ret != 0: -+ with self.session_lock: -+ self.session_failures[ -+ req.mooncake_session_id -+ ] += 1 ++ remote_addr = NetworkAddress( ++ req.endpoint, req.dst_port ++ ).to_host_port_str() + self.record_failure( + kv_chunk.room, -+ f"Failed to send extra state chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}", ++ f"Failed to send extra state chunk of {kv_chunk.room} to " ++ f"{remote_addr}", + ) + self.update_status(kv_chunk.room, KVPoll.Failed) + self.sync_status_to_decode_endpoint( @@ -375,77 +334,26 @@ index d0d4efd958..b3a207063e 100644 + req.dst_port, + req.room, + KVPoll.Failed, -+ local_rank, ++ prefill_unique_rank, + ) + break # Only the last chunk we need to send the aux data ret = self.send_aux( -@@ -895,6 +929,11 @@ class MooncakeKVManager(CommonKVManager): - target_rank_registration_info.dst_aux_ptrs, - ) - polls.append(True if ret == 0 else False) -+ if ret != 0: -+ # Mark session as failed to avoid hanging -+ # on subsequent batch_transfer_sync calls -+ with self.session_lock: -+ self.session_failures[req.mooncake_session_id] += 1 - dst_ranks_infos.append( - (req.endpoint, req.dst_port, req.room) - ) -@@ -977,15 +1016,20 @@ class MooncakeKVManager(CommonKVManager): - - if status == KVPoll.Success: - if bootstrap_room in self.request_status: -- self.prefill_response_tracker[bootstrap_room].add(prefill_rank) -+ # Guard against TOCTOU race: clear() may remove the entry -+ # between the request_status check and dict access here. - expected_response_num = ( -- self.required_prefill_response_num_table[bootstrap_room] -+ self.required_prefill_response_num_table.get(bootstrap_room) - ) -- arrived_response_num = len( -- self.prefill_response_tracker[bootstrap_room] -- ) -- if arrived_response_num == expected_response_num: -- self.update_status(bootstrap_room, KVPoll.Success) -+ if expected_response_num is not None: -+ self.prefill_response_tracker[bootstrap_room].add( -+ prefill_rank -+ ) -+ arrived_response_num = len( -+ self.prefill_response_tracker[bootstrap_room] -+ ) -+ if arrived_response_num == expected_response_num: -+ self.update_status(bootstrap_room, KVPoll.Success) - elif status == KVPoll.Failed: - self.record_failure( - bootstrap_room, -@@ -1266,7 +1310,10 @@ class MooncakeKVReceiver(CommonKVReceiver): - super().__init__(mgr, bootstrap_addr, bootstrap_room, prefill_dp_rank) - - self.kv_mgr.addr_to_rooms_tracker[self.bootstrap_addr].add(self.bootstrap_room) -- self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) -+ # Only transition to WaitingForInput if bootstrap succeeded; -+ # if super().__init__() set status to Failed, do not override it. -+ if self.bootstrap_infos is not None: -+ self.kv_mgr.update_status(self.bootstrap_room, KVPoll.WaitingForInput) - - def _register_kv_args(self): - for bootstrap_info in self.bootstrap_infos: diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py -index fbc8016351..7de53ba292 100644 +index 8eadf81954..c180ce79f3 100644 --- a/python/sglang/srt/disaggregation/prefill.py +++ b/python/sglang/srt/disaggregation/prefill.py -@@ -20,6 +20,7 @@ Life cycle of a request in the prefill server +@@ -20,6 +20,8 @@ Life cycle of a request in the prefill server from __future__ import annotations import logging +import os - import time ++import time from collections import deque from http import HTTPStatus -@@ -167,6 +168,7 @@ class PrefillBootstrapQueue: + from typing import TYPE_CHECKING, List, Optional +@@ -165,6 +167,7 @@ class PrefillBootstrapQueue: kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = ( self.metadata_buffers.get_buf_infos() ) @@ -453,11 +361,10 @@ index fbc8016351..7de53ba292 100644 kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device kv_args.gpu_id = self.scheduler.gpu_id -@@ -276,6 +278,12 @@ class PrefillBootstrapQueue: - [req.disagg_kv_sender for req in self.queue], self.gloo_group +@@ -290,6 +293,11 @@ class PrefillBootstrapQueue: + self.scheduler.attn_tp_cpu_group, ) -+ # Bootstrap timeout: if a request has been stuck in Bootstrapping for too long, treat it as failed. + bootstrap_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) @@ -466,11 +373,10 @@ index fbc8016351..7de53ba292 100644 for i, (req, poll) in enumerate(zip(self.queue, polls)): if rids_to_check is not None: # if req not in reqs_info_to_check, skip -@@ -283,6 +291,27 @@ class PrefillBootstrapQueue: +@@ -297,6 +305,26 @@ class PrefillBootstrapQueue: continue if poll == KVPoll.Bootstrapping: -+ # Check for bootstrap timeout + entry_time = getattr( + req.time_stats, + "prefill_bootstrap_queue_entry_time", @@ -494,7 +400,7 @@ index fbc8016351..7de53ba292 100644 continue elif poll == KVPoll.Failed: error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}" -@@ -335,6 +364,15 @@ class PrefillBootstrapQueue: +@@ -346,6 +374,15 @@ class PrefillBootstrapQueue: else: return bootstrapped_reqs, failed_reqs @@ -510,31 +416,34 @@ index fbc8016351..7de53ba292 100644 class SchedulerDisaggregationPrefillMixin: """ -@@ -547,6 +585,18 @@ class SchedulerDisaggregationPrefillMixin: - - self.maybe_send_health_check_signal() - +@@ -568,12 +605,17 @@ class SchedulerDisaggregationPrefillMixin: + self.send_kv_chunk(req, last_chunk=False, end_idx=req.tmp_end_idx) + req.time_stats.set_last_chunked_prefill_finish_time() + +- can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) +- self.report_prefill_stats( +- prefill_stats=batch.prefill_stats, +- can_run_cuda_graph=can_run_cuda_graph, +- dp_cooperation_info=batch.dp_cooperation_info, +- ) + if ( + self.current_scheduler_metrics_enabled + and hasattr(batch, "prefill_stats") + and batch.prefill_stats is not None + ): + can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) -+ self.log_prefill_stats( ++ self.report_prefill_stats( + prefill_stats=batch.prefill_stats, + can_run_cuda_graph=can_run_cuda_graph, + dp_cooperation_info=getattr(batch, "dp_cooperation_info", None), + ) -+ + def process_disagg_prefill_inflight_queue( self: Scheduler, rids_to_check: Optional[List[str]] = None - ) -> List[Req]: -@@ -564,6 +614,13 @@ class SchedulerDisaggregationPrefillMixin: +@@ -593,6 +635,11 @@ class SchedulerDisaggregationPrefillMixin: self.attn_tp_cpu_group, ) -+ # Transfer timeout: if a request has been in the inflight queue for too long -+ # (e.g., stuck in WaitingForInput/Transferring), treat it as failed. + transfer_timeout = float( + os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600") + ) @@ -543,18 +452,11 @@ index fbc8016351..7de53ba292 100644 undone_reqs: List[Req] = [] # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue for req, poll in zip(self.disagg_prefill_inflight_queue, polls): -@@ -573,10 +630,35 @@ class SchedulerDisaggregationPrefillMixin: - undone_reqs.append(req) +@@ -618,7 +665,29 @@ class SchedulerDisaggregationPrefillMixin: continue -- assert poll == KVPoll.Success or poll == KVPoll.Failed -+ if poll not in (KVPoll.Success, KVPoll.Failed): -+ undone_reqs.append(req) -+ continue - if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]: - undone_reqs.append(req) -+ # Check for transfer timeout + entry_time = getattr( + req.time_stats, + "prefill_transfer_queue_entry_time", @@ -567,7 +469,7 @@ index fbc8016351..7de53ba292 100644 + f"{req.rid=} {req.bootstrap_room=}" + ) + logger.error(error_message) -+ release_kv_cache(req, self.tree_cache) # unlock the tree ++ release_kv_cache(req, self.tree_cache) + prepare_abort( + req, error_message, status_code=HTTPStatus.GATEWAY_TIMEOUT + ) @@ -582,10 +484,10 @@ index fbc8016351..7de53ba292 100644 release_kv_cache(req, self.tree_cache) # unlock the tree req.finished_reason = FINISH_LENGTH(length=0) diff --git a/python/sglang/srt/disaggregation/utils.py b/python/sglang/srt/disaggregation/utils.py -index 6d58f415a7..84723c342c 100644 +index d7956a6048..0ced278713 100644 --- a/python/sglang/srt/disaggregation/utils.py +++ b/python/sglang/srt/disaggregation/utils.py -@@ -21,6 +21,17 @@ if TYPE_CHECKING: +@@ -28,6 +28,17 @@ if TYPE_CHECKING: # Constants & Enums ######################### FAKE_BOOTSTRAP_HOST = "2.2.2.2" @@ -603,9 +505,9 @@ index 6d58f415a7..84723c342c 100644 class DisaggregationMode(Enum): -@@ -139,46 +150,35 @@ class MetadataBuffers: +@@ -193,46 +204,35 @@ class MetadataBuffers: self.bootstrap_room = torch.zeros( - (size, 8), dtype=torch.uint64, device=device + (size, 8), dtype=bootstrap_room_dtype, device=device ) + # Prefill-side PD timing (8 floats, padded to 16 for RDMA alignment). + # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap, @@ -675,7 +577,7 @@ index 6d58f415a7..84723c342c 100644 def get_buf(self, idx: int): return ( self.output_ids[idx], -@@ -191,8 +191,12 @@ class MetadataBuffers: +@@ -245,8 +245,12 @@ class MetadataBuffers: self.output_topk_index[idx], self.output_hidden_states[idx], self.bootstrap_room[idx], @@ -688,7 +590,7 @@ index 6d58f415a7..84723c342c 100644 def set_buf(self, req: Req): self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0] -@@ -237,6 +241,84 @@ class MetadataBuffers: +@@ -294,6 +298,99 @@ class MetadataBuffers: self.bootstrap_room[req.metadata_buffer_index, 0] = ( req.bootstrap_room if req.bootstrap_room is not None else 0 ) @@ -731,12 +633,27 @@ index 6d58f415a7..84723c342c 100644 + else 0.0 + ) + ++ bootstrap_duration = 0.0 ++ alloc_waiting_duration = 0.0 ++ if ( ++ time_stats.prefill_bootstrap_queue_entry_time > 0 ++ and time_stats.bootstrap_done_time > 0 ++ ): ++ bootstrap_duration = ( ++ time_stats.bootstrap_done_time ++ - time_stats.prefill_bootstrap_queue_entry_time ++ ) ++ if time_stats.bootstrap_done_time > 0 and time_stats.wait_queue_entry_time > 0: ++ alloc_waiting_duration = ( ++ time_stats.wait_queue_entry_time - time_stats.bootstrap_done_time ++ ) ++ + return ( + bootstrap_queue_duration, + prefill_forward_duration, + 0.0, -+ max(0.0, time_stats.bootstrap_duration), -+ max(0.0, time_stats.alloc_waiting_duration), ++ max(0.0, bootstrap_duration), ++ max(0.0, alloc_waiting_duration), + max(0.0, time_stats.transfer_speed_gb_s), + max(0.0, time_stats.transfer_total_mb), + float(max(0, time_stats.prefill_retry_count)), @@ -774,10 +691,10 @@ index 6d58f415a7..84723c342c 100644 ######################### diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py -index 0ed5a1b44b..67e33c650d 100644 +index d864e4abaa..3a000a80f2 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py -@@ -52,6 +52,7 @@ from sglang.srt.managers.io_struct import ( +@@ -69,6 +69,7 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterReqInput, MultimodalDataInputFormat, OpenSessionReqInput, @@ -785,7 +702,7 @@ index 0ed5a1b44b..67e33c650d 100644 ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, RpcReqInput, -@@ -641,6 +642,24 @@ class Engine(EngineBase): +@@ -957,6 +958,24 @@ class Engine(EngineScoreMixin, EngineBase): self.tokenizer_manager.update_weights_from_ipc(obj, None) ) @@ -811,10 +728,10 @@ index 0ed5a1b44b..67e33c650d 100644 """Get weights by parameter name.""" obj = GetWeightsByNameReqInput(name=name, truncate_size=truncate_size) diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py -index 1d6816c010..402b42e05b 100644 +index 6978e0c062..80dc159e8f 100644 --- a/python/sglang/srt/entrypoints/http_server.py +++ b/python/sglang/srt/entrypoints/http_server.py -@@ -115,6 +115,7 @@ from sglang.srt.managers.io_struct import ( +@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import ( OpenSessionReqInput, ParseFunctionCallReq, PauseGenerationReqInput, @@ -822,7 +739,7 @@ index 1d6816c010..402b42e05b 100644 ProfileReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -574,10 +575,8 @@ async def model_info(): +@@ -582,10 +583,8 @@ async def model_info(): @app.get("/weight_version") async def weight_version(): """Get the current weight version.""" @@ -835,7 +752,7 @@ index 1d6816c010..402b42e05b 100644 @app.get("/get_server_info") -@@ -594,9 +593,19 @@ async def get_server_info(): +@@ -602,9 +601,19 @@ async def get_server_info(): async def server_info(): """Get the server information.""" # Returns internal states per DP. @@ -858,7 +775,7 @@ index 1d6816c010..402b42e05b 100644 # This field is not serializable. if hasattr(_global_state.tokenizer_manager.server_args, "model_config"): -@@ -1084,6 +1093,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re +@@ -1121,6 +1130,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST) @@ -883,39 +800,19 @@ index 1d6816c010..402b42e05b 100644 @auth_level(AuthLevel.ADMIN_OPTIONAL) async def update_weight_version(obj: UpdateWeightVersionReqInput, request: Request): diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py -index 8293796a2e..bff34e4221 100644 +index dfc5507de0..be9501b05a 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py -@@ -244,6 +244,7 @@ class Envs: +@@ -242,6 +242,7 @@ class Envs: SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2) SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300) SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX") + SLIME_ENABLE_PROFILING = EnvBool(False) - - # Scheduler: others: - SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period. -diff --git a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py -index 1cdf65b91c..4783cd18fb 100644 ---- a/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py -+++ b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py -@@ -630,7 +630,6 @@ def _get_k_and_s_triton( - page_indices, - k_out, - s_out, -- seq_len, - page_size, - buf_numel_per_page, - index_head_dim, -@@ -647,7 +646,6 @@ def _get_k_and_s_triton_kernel( - page_indices_ptr, - k_out_ptr, - s_out_ptr, -- seq_len: tl.constexpr, - page_size: tl.constexpr, - buf_numel_per_page: tl.constexpr, - index_head_dim: tl.constexpr, + SGLANG_DISAGGREGATION_ALL_CP_RANKS_TRANSFER = EnvBool(False) + # Extra slots in req_to_token_pool for decode workers (only effective when + # max_num_reqs > 32). Increases pool capacity so more KV cache transfers diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py -index ca54a931b7..3540f77bae 100644 +index 02ef4e2440..fd5a43cce8 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,6 +1,7 @@ @@ -926,7 +823,7 @@ index ca54a931b7..3540f77bae 100644 from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple -@@ -201,14 +202,31 @@ class Indexer(MultiPlatformOp): +@@ -213,14 +214,31 @@ class Indexer(MultiPlatformOp): prefix=add_prefix("weights_proj", prefix), ) self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32) @@ -960,10 +857,10 @@ index ca54a931b7..3540f77bae 100644 ) self.block_size = block_size self.scale_fmt = scale_fmt -@@ -244,6 +262,11 @@ class Indexer(MultiPlatformOp): - x = x.to(self.weights_proj.weight.dtype) - weights, _ = self.weights_proj(x) - weights = weights.float() +@@ -266,6 +284,11 @@ class Indexer(MultiPlatformOp): + @torch.compile(dynamic=True) if not _is_hip else lambda f: f + def _get_logits_head_gate(self, x: torch.Tensor, q_scale: torch.Tensor): + weights = self._weights_proj_bf16_in_fp32_out(x) + if weights.shape[1] < q_scale.shape[1]: + assert q_scale.shape[1] % weights.shape[1] == 0 + weights = weights.repeat_interleave( @@ -972,7 +869,7 @@ index ca54a931b7..3540f77bae 100644 weights = weights * self.n_heads**-0.5 weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale return weights -@@ -982,15 +1005,26 @@ class Indexer(MultiPlatformOp): +@@ -1078,6 +1101,9 @@ class Indexer(MultiPlatformOp): query, key = self._get_q_k_bf16( q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch ) @@ -981,15 +878,8 @@ index ca54a931b7..3540f77bae 100644 + query = query.repeat_interleave(32 // query.shape[1], dim=1) q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt) with torch.cuda.stream(self.alt_stream): - k_fp8, k_scale = act_quant(key, self.block_size, self.scale_fmt) - current_stream.wait_stream(self.alt_stream) -+ if weights.shape[1] < q_scale.shape[1]: -+ assert q_scale.shape[1] % weights.shape[1] == 0 -+ weights = weights.repeat_interleave( -+ q_scale.shape[1] // weights.shape[1], dim=1 -+ ) - weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale - else: + self._store_index_k_cache( +@@ -1092,6 +1118,9 @@ class Indexer(MultiPlatformOp): query, key = self._get_q_k_bf16( q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch ) @@ -1000,10 +890,10 @@ index ca54a931b7..3540f77bae 100644 if enable_dual_stream: current_stream = torch.cuda.current_stream() diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -index de8a07ab30..5c9f4813a6 100644 +index 72483f4ea6..2e1148d189 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py -@@ -697,6 +697,7 @@ class FusedMoE(torch.nn.Module): +@@ -702,6 +702,7 @@ class FusedMoE(torch.nn.Module): "CompressedTensorsWNA16TritonMoE", ] ) @@ -1011,7 +901,7 @@ index de8a07ab30..5c9f4813a6 100644 else loaded_weight ) -@@ -916,6 +917,7 @@ class FusedMoE(torch.nn.Module): +@@ -921,6 +922,7 @@ class FusedMoE(torch.nn.Module): "CompressedTensorsWNA16TritonMoE", ] ) @@ -1081,10 +971,10 @@ index 00bd687555..12d5577af2 100644 def get_routed_experts( diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -index 4cbfed6f90..88b4527443 100644 +index a13c53af4d..1d80d06b13 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py -@@ -499,7 +499,7 @@ class CompressedTensorsConfig(QuantizationConfig): +@@ -500,7 +500,7 @@ class CompressedTensorsConfig(QuantizationConfig): ) is_static = not weight_quant.dynamic @@ -1093,7 +983,7 @@ index 4cbfed6f90..88b4527443 100644 def _is_mxint4a16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool: input_quant_none = input_quant is None -@@ -968,6 +968,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): +@@ -969,6 +969,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase): def process_weights_after_loading(self, layer: torch.nn.Module) -> None: layer.scheme.process_weights_after_loading(layer) @@ -1104,7 +994,7 @@ index 4cbfed6f90..88b4527443 100644 self, layer: torch.nn.Module, diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py -index 6264f36d04..f0310e305e 100644 +index 7a8fb65421..f1c85899cd 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -17,7 +17,10 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( @@ -1170,25 +1060,18 @@ index 6264f36d04..f0310e305e 100644 w13_g_idx = torch.nn.Parameter( torch.empty( num_experts, -@@ -225,11 +254,14 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): +@@ -231,6 +260,10 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): + layer._original_shapes["w2_weight_scale"] = tuple(w2_scale.shape) + layer._original_shapes["w13_weight_scale"] = tuple(w13_scale.shape) - # Force record: these are the target GPTQ shapes for rollback. - layer._original_shapes["w13_weight_packed"] = tuple(w13_weight.shape) -- layer._original_shapes["w2_weight_packed"] = tuple(w2_weight.shape) -+ layer._original_shapes["w13_weight_scale"] = tuple(w13_scale.shape) + if not self.sym: + layer._original_shapes["w13_weight_zero_point"] = w13_qzeros.shape - -- # Also record the shapes of the scales. -+ layer._original_shapes["w2_weight_packed"] = tuple(w2_weight.shape) - layer._original_shapes["w2_weight_scale"] = tuple(w2_scale.shape) -- layer._original_shapes["w13_weight_scale"] = tuple(w13_scale.shape) -+ if not self.sym: + layer._original_shapes["w2_weight_zero_point"] = tuple(w2_qzeros.shape) - ++ def process_weights_after_loading(self, layer: torch.nn.Module) -> None: -@@ -334,6 +366,24 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): + # Skip if the layer is already converted to Marlin format to prevent double-packing. +@@ -334,6 +367,24 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): ) replace_tensor("w2_weight_scale", marlin_w2_scales) @@ -1213,7 +1096,7 @@ index 6264f36d04..f0310e305e 100644 layer.is_marlin_converted = True def restore_weights_before_loading(self, layer: torch.nn.Module): -@@ -399,6 +449,8 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): +@@ -399,6 +450,8 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme): g_idx2=layer.w2_weight_g_idx, sort_indices1=layer.w13_g_idx_sort_indices, sort_indices2=layer.w2_g_idx_sort_indices, @@ -1222,84 +1105,17 @@ index 6264f36d04..f0310e305e 100644 num_bits=self.num_bits, is_k_full=self.is_k_full, routed_scaling_factor=self.moe_runner_config.routed_scaling_factor, -diff --git a/python/sglang/srt/managers/detokenizer_manager.py b/python/sglang/srt/managers/detokenizer_manager.py -index 6522278603..7d3a5d0c4c 100644 ---- a/python/sglang/srt/managers/detokenizer_manager.py -+++ b/python/sglang/srt/managers/detokenizer_manager.py -@@ -405,6 +405,17 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin): - prefill_launch_delay=recv_obj.prefill_launch_delay, - prefill_launch_latency=recv_obj.prefill_launch_latency, - prefill_finished_ts=recv_obj.prefill_finished_ts, -+ pd_prefill_bootstrap_queue_duration=recv_obj.pd_prefill_bootstrap_queue_duration, -+ pd_prefill_forward_duration=recv_obj.pd_prefill_forward_duration, -+ pd_prefill_transfer_queue_duration=recv_obj.pd_prefill_transfer_queue_duration, -+ pd_decode_prealloc_duration=recv_obj.pd_decode_prealloc_duration, -+ pd_decode_transfer_duration=recv_obj.pd_decode_transfer_duration, -+ pd_decode_forward_duration=recv_obj.pd_decode_forward_duration, -+ pd_bootstrap_duration=recv_obj.pd_bootstrap_duration, -+ pd_alloc_waiting_duration=recv_obj.pd_alloc_waiting_duration, -+ pd_transfer_speed_gb_s=recv_obj.pd_transfer_speed_gb_s, -+ pd_transfer_total_mb=recv_obj.pd_transfer_total_mb, -+ pd_prefill_retry_count=recv_obj.pd_prefill_retry_count, - ) - - def handle_multimodal_decode_req(self, recv_obj: BatchMultimodalDecodeReq): diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index ff17745673..f947e71d7d 100644 +index bd97965345..e6a147c1b4 100644 --- a/python/sglang/srt/managers/io_struct.py +++ b/python/sglang/srt/managers/io_struct.py -@@ -101,6 +101,42 @@ class RequestTimingMetricsMixin: - # This marks when the prefill computation finishes. - prefill_finished_ts: Optional[List[Optional[float]]] - -+ # --- PD disaggregation timing fields --- -+ # All fields are None when profiling is disabled or not in PD disaggregation mode. -+ -+ # P instance: duration spent in bootstrap queue before entering the wait queue. -+ pd_prefill_bootstrap_queue_duration: Optional[List[Optional[float]]] -+ -+ # P instance: duration for the actual prefill forward computation. -+ pd_prefill_forward_duration: Optional[List[Optional[float]]] -+ -+ # P instance: duration spent in the KV transfer queue. -+ pd_prefill_transfer_queue_duration: Optional[List[Optional[float]]] -+ -+ # D instance: duration waiting for KV cache slot pre-allocation. -+ pd_decode_prealloc_duration: Optional[List[Optional[float]]] -+ -+ # D instance: duration waiting for the KV cache transfer to complete. -+ pd_decode_transfer_duration: Optional[List[Optional[float]]] -+ -+ # D instance: duration for the actual decode forward computation. -+ pd_decode_forward_duration: Optional[List[Optional[float]]] -+ -+ # Bootstrap handshake duration (P and D instances). -+ pd_bootstrap_duration: Optional[List[Optional[float]]] -+ -+ # KV cache allocation waiting duration (P and D instances). -+ pd_alloc_waiting_duration: Optional[List[Optional[float]]] -+ -+ # KV cache transfer speed in GB/s. -+ pd_transfer_speed_gb_s: Optional[List[Optional[float]]] -+ -+ # Total KV cache transferred in MB. -+ pd_transfer_total_mb: Optional[List[Optional[float]]] -+ -+ # Number of prefill retries (P instance only). -+ pd_prefill_retry_count: Optional[List[Optional[int]]] -+ - - @dataclass - class SpeculativeDecodingMetricsMixin: -@@ -1403,6 +1439,20 @@ class UpdateWeightsFromIPCReqOutput(BaseReq): - message: str +@@ -1449,6 +1449,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq): + pass +@dataclass +class PostProcessWeightsReqInput(BaseReq): -+ # Whether to restore weights before loading new weights + restore_weights_before_load: bool = False -+ # Whether to enable quantization post-processing + post_process_quantization: bool = False + + @@ -1310,186 +1126,32 @@ index ff17745673..f947e71d7d 100644 + + @dataclass - class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq): - success: bool -@@ -1802,6 +1852,10 @@ class GetLoadReqOutput(BaseReq): + class CheckWeightsReqInput(BaseReq): + action: str +@@ -1753,6 +1765,8 @@ class GetLoadReqOutput(BaseReq): num_waiting_reqs: int num_tokens: int ts_tic: float -+ # Per-queue breakdown: list of {name, num_reqs, num_tokens, reqs: [{rid, seqlen, input_len, output_len}]} + queue_details: Optional[List[Dict[str, Any]]] = None -+ # Running batch info + running_details: Optional[Dict[str, Any]] = None @dataclass diff --git a/python/sglang/srt/managers/multi_tokenizer_mixin.py b/python/sglang/srt/managers/multi_tokenizer_mixin.py -index e1236aa0f3..daa598a1f6 100644 +index e0a1669fb3..fbbb6bb12b 100644 --- a/python/sglang/srt/managers/multi_tokenizer_mixin.py +++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py -@@ -142,6 +142,39 @@ def _handle_output_by_index(output, i): - prefill_finished_ts=_extract_field_by_index( - output, "prefill_finished_ts", i - ), -+ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_bootstrap_queue_duration", i -+ ), -+ pd_prefill_forward_duration=_extract_field_by_index( -+ output, "pd_prefill_forward_duration", i -+ ), -+ pd_prefill_transfer_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_transfer_queue_duration", i -+ ), -+ pd_decode_prealloc_duration=_extract_field_by_index( -+ output, "pd_decode_prealloc_duration", i -+ ), -+ pd_decode_transfer_duration=_extract_field_by_index( -+ output, "pd_decode_transfer_duration", i -+ ), -+ pd_decode_forward_duration=_extract_field_by_index( -+ output, "pd_decode_forward_duration", i -+ ), -+ pd_bootstrap_duration=_extract_field_by_index( -+ output, "pd_bootstrap_duration", i -+ ), -+ pd_alloc_waiting_duration=_extract_field_by_index( -+ output, "pd_alloc_waiting_duration", i -+ ), -+ pd_transfer_speed_gb_s=_extract_field_by_index( -+ output, "pd_transfer_speed_gb_s", i -+ ), -+ pd_transfer_total_mb=_extract_field_by_index( -+ output, "pd_transfer_total_mb", i -+ ), -+ pd_prefill_retry_count=_extract_field_by_index( -+ output, "pd_prefill_retry_count", i -+ ), - finished_reasons=_extract_field_by_index(output, "finished_reasons", i), - decoded_texts=_extract_field_by_index(output, "decoded_texts", i), - decode_ids=_extract_field_by_index(output, "decode_ids", i), -@@ -211,6 +244,50 @@ def _handle_output_by_index(output, i): - elif isinstance(output, BatchEmbeddingOutput): - new_output = BatchEmbeddingOutput( - rids=[output.rids[i]], -+ queue_time=_extract_field_by_index(output, "queue_time", i), -+ forward_entry_time=_extract_field_by_index(output, "forward_entry_time", i), -+ prefill_launch_delay=_extract_field_by_index( -+ output, "prefill_launch_delay", i -+ ), -+ prefill_launch_latency=_extract_field_by_index( -+ output, "prefill_launch_latency", i -+ ), -+ prefill_finished_ts=_extract_field_by_index( -+ output, "prefill_finished_ts", i -+ ), -+ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_bootstrap_queue_duration", i -+ ), -+ pd_prefill_forward_duration=_extract_field_by_index( -+ output, "pd_prefill_forward_duration", i -+ ), -+ pd_prefill_transfer_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_transfer_queue_duration", i -+ ), -+ pd_decode_prealloc_duration=_extract_field_by_index( -+ output, "pd_decode_prealloc_duration", i -+ ), -+ pd_decode_transfer_duration=_extract_field_by_index( -+ output, "pd_decode_transfer_duration", i -+ ), -+ pd_decode_forward_duration=_extract_field_by_index( -+ output, "pd_decode_forward_duration", i -+ ), -+ pd_bootstrap_duration=_extract_field_by_index( -+ output, "pd_bootstrap_duration", i -+ ), -+ pd_alloc_waiting_duration=_extract_field_by_index( -+ output, "pd_alloc_waiting_duration", i -+ ), -+ pd_transfer_speed_gb_s=_extract_field_by_index( -+ output, "pd_transfer_speed_gb_s", i -+ ), -+ pd_transfer_total_mb=_extract_field_by_index( -+ output, "pd_transfer_total_mb", i -+ ), -+ pd_prefill_retry_count=_extract_field_by_index( -+ output, "pd_prefill_retry_count", i -+ ), - finished_reasons=_extract_field_by_index(output, "finished_reasons", i), - embeddings=_extract_field_by_index(output, "embeddings", i), - prompt_tokens=_extract_field_by_index(output, "prompt_tokens", i), -@@ -239,6 +316,39 @@ def _handle_output_by_index(output, i): - prefill_finished_ts=_extract_field_by_index( - output, "prefill_finished_ts", i - ), -+ pd_prefill_bootstrap_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_bootstrap_queue_duration", i -+ ), -+ pd_prefill_forward_duration=_extract_field_by_index( -+ output, "pd_prefill_forward_duration", i -+ ), -+ pd_prefill_transfer_queue_duration=_extract_field_by_index( -+ output, "pd_prefill_transfer_queue_duration", i -+ ), -+ pd_decode_prealloc_duration=_extract_field_by_index( -+ output, "pd_decode_prealloc_duration", i -+ ), -+ pd_decode_transfer_duration=_extract_field_by_index( -+ output, "pd_decode_transfer_duration", i -+ ), -+ pd_decode_forward_duration=_extract_field_by_index( -+ output, "pd_decode_forward_duration", i -+ ), -+ pd_bootstrap_duration=_extract_field_by_index( -+ output, "pd_bootstrap_duration", i -+ ), -+ pd_alloc_waiting_duration=_extract_field_by_index( -+ output, "pd_alloc_waiting_duration", i -+ ), -+ pd_transfer_speed_gb_s=_extract_field_by_index( -+ output, "pd_transfer_speed_gb_s", i -+ ), -+ pd_transfer_total_mb=_extract_field_by_index( -+ output, "pd_transfer_total_mb", i -+ ), -+ pd_prefill_retry_count=_extract_field_by_index( -+ output, "pd_prefill_retry_count", i -+ ), - finished_reasons=_extract_field_by_index(output, "finished_reasons", i), - output_strs=_extract_field_by_index(output, "output_strs", i), - output_ids=_extract_field_by_index(output, "output_ids", i), -@@ -524,6 +634,60 @@ def monkey_patch_uvicorn_multiprocessing(timeout: float = 10): +@@ -496,6 +496,35 @@ def monkey_patch_uvicorn_multiprocessing(timeout: float = 10): "uvicorn.supervisors.multiprocess not found, skipping monkey patch" ) -+ # Fix stdin fd issue when running under Ray (or other managed -+ # environments where stdin may not be a real terminal): -+ # -+ # Uvicorn's get_subprocess() captures sys.stdin.fileno() in the parent -+ # and passes it to spawn'd children, which call os.fdopen(stdin_fileno) -+ # to re-attach stdin. This is intended for interactive debugging (e.g. -+ # pdb attach to a child worker). -+ # -+ # In Ray Actors, sys.stdin.fileno() succeeds in the parent (returns a -+ # valid fd number), but the fd is not inheritable across spawn. The -+ # child's os.fdopen() then crashes with OSError: [Errno 9] Bad file -+ # descriptor, killing every tokenizer worker. -+ # -+ # Instead of unconditionally disabling stdin passthrough, we probe -+ # whether the fd is truly usable by dup'ing it. If os.dup() fails, -+ # the fd won't survive spawn either, so we fall back to None. In a -+ # normal terminal environment os.dup() succeeds and debugging ability -+ # is preserved. + try: -+ import uvicorn._subprocess as _uv_sub -+ import uvicorn.supervisors.multiprocess as _uv_mp ++ import uvicorn._subprocess as uvicorn_subprocess ++ import uvicorn.supervisors.multiprocess as uvicorn_multiprocess + + def _safe_get_stdin_fileno(): -+ """Return stdin fileno only if it is genuinely usable.""" + try: + fileno = sys.stdin.fileno() -+ # Verify the fd is valid and duplicable — if it isn't, -+ # spawn'd children won't be able to reopen it either. + dup_fd = os.dup(fileno) + os.close(dup_fd) + return fileno @@ -1497,22 +1159,18 @@ index e1236aa0f3..daa598a1f6 100644 + return None + + def _patched_get_subprocess(config, target, sockets): -+ stdin_fileno = _safe_get_stdin_fileno() + kwargs = { + "config": config, + "target": target, + "sockets": sockets, -+ "stdin_fileno": stdin_fileno, ++ "stdin_fileno": _safe_get_stdin_fileno(), + } -+ return _uv_sub.spawn.Process( -+ target=_uv_sub.subprocess_started, kwargs=kwargs ++ return uvicorn_subprocess.spawn.Process( ++ target=uvicorn_subprocess.subprocess_started, kwargs=kwargs + ) + -+ # Must patch both: the supervisor module caches its own reference -+ # to get_subprocess at import time via -+ # ``from uvicorn._subprocess import get_subprocess``. -+ _uv_sub.get_subprocess = _patched_get_subprocess -+ _uv_mp.get_subprocess = _patched_get_subprocess ++ uvicorn_subprocess.get_subprocess = _patched_get_subprocess ++ uvicorn_multiprocess.get_subprocess = _patched_get_subprocess + except Exception: + pass + @@ -1520,10 +1178,10 @@ index e1236aa0f3..daa598a1f6 100644 class SenderWrapper: def __init__(self, port_args: PortArgs, send_to_scheduler: zmq.Socket): diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py -index c079957980..dd8ca7167d 100644 +index 0b26be6c6d..2ea1042cf9 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py -@@ -1869,7 +1869,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): +@@ -1972,7 +1972,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): while first_iter or ( not self.check_decode_mem(selected_indices=sorted_indices) ): @@ -1536,18 +1194,18 @@ index c079957980..dd8ca7167d 100644 break diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index a9ff0ac94b..a50dd5122b 100644 +index 67af2d0de9..122ddb3874 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py -@@ -114,6 +114,7 @@ from sglang.srt.managers.io_struct import ( +@@ -120,6 +120,7 @@ from sglang.srt.managers.io_struct import ( + LoadLoRAAdapterReqOutput, OpenSessionReqInput, - OpenSessionReqOutput, PauseGenerationReqInput, + PostProcessWeightsReqInput, ProfileReq, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, -@@ -1063,6 +1064,7 @@ class Scheduler( +@@ -1232,6 +1233,7 @@ class Scheduler( ), (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor), (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc), @@ -1555,239 +1213,24 @@ index a9ff0ac94b..a50dd5122b 100644 (GetWeightsByNameReqInput, self.get_weights_by_name), (ReleaseMemoryOccupationReqInput, self.release_memory_occupation), (ResumeMemoryOccupationReqInput, self.resume_memory_occupation), -diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py -index 30b2732b9f..68090b1617 100644 ---- a/python/sglang/srt/managers/scheduler_metrics_mixin.py -+++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py -@@ -609,12 +609,54 @@ class SchedulerMetricsMixin: - num_tokens += sum(req.seqlen for queue in waiting_queues for req in queue) - num_waiting_reqs = sum(len(queue) for queue in waiting_queues) - -+ # Collect per-queue details -+ queue_names = ["waiting_queue"] -+ if self.disaggregation_mode == DisaggregationMode.PREFILL: -+ queue_names.append("bootstrap_queue") -+ elif self.disaggregation_mode == DisaggregationMode.DECODE: -+ queue_names.append("prealloc_queue") -+ queue_names.append("transfer_queue") -+ queue_names.append("retracted_queue") -+ -+ queue_details = [] -+ for name, queue in zip(queue_names, waiting_queues): -+ reqs_info = [] -+ for req in queue: -+ reqs_info.append( -+ { -+ "seqlen": req.seqlen, -+ } -+ ) -+ queue_details.append( -+ { -+ "name": name, -+ "num_reqs": len(queue), -+ "num_tokens": sum(r["seqlen"] for r in reqs_info), -+ "reqs": reqs_info, -+ } -+ ) -+ -+ # Collect running batch details -+ running_reqs_info = [] -+ for req in self.running_batch.reqs: -+ running_reqs_info.append( -+ { -+ "seqlen": req.seqlen, -+ } -+ ) -+ running_details = { -+ "num_reqs": len(self.running_batch.reqs), -+ "reqs": running_reqs_info, -+ } -+ - return GetLoadReqOutput( - dp_rank=self.dp_rank, - num_reqs=len(self.running_batch.reqs) + num_waiting_reqs, - num_waiting_reqs=num_waiting_reqs, - num_tokens=num_tokens, - ts_tic=time.perf_counter(), -+ queue_details=queue_details, -+ running_details=running_details, - ) - - def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -index 482bc6ca66..fbc4864176 100644 +index 496cd96656..cf2d43015a 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py -@@ -922,6 +922,18 @@ class SchedulerOutputProcessorMixin: - prefill_launch_delays = [] - prefill_launch_latencies = [] - prefill_finished_timestamps = [] -+ profiling_enabled = envs.SLIME_ENABLE_PROFILING.get() -+ pd_prefill_bootstrap_queue_durations = [] if profiling_enabled else None -+ pd_prefill_forward_durations = [] if profiling_enabled else None -+ pd_prefill_transfer_queue_durations = [] if profiling_enabled else None -+ pd_decode_prealloc_durations = [] if profiling_enabled else None -+ pd_decode_transfer_durations = [] if profiling_enabled else None -+ pd_decode_forward_durations = [] if profiling_enabled else None -+ pd_bootstrap_durations = [] if profiling_enabled else None -+ pd_alloc_waiting_durations = [] if profiling_enabled else None -+ pd_transfer_speeds_gb_s = [] if profiling_enabled else None -+ pd_transfer_totals_mb = [] if profiling_enabled else None -+ pd_prefill_retry_counts = [] if profiling_enabled else None - - if return_logprob: - input_token_logprobs_val = [] -@@ -1037,6 +1049,40 @@ class SchedulerOutputProcessorMixin: - prefill_finished_timestamps.append( - req.time_stats.get_prefill_finished_ts() - ) -+ if profiling_enabled: -+ pd_prefill_bootstrap_queue_durations.append( -+ req.time_stats.get_pd_prefill_bootstrap_queue_duration() -+ ) -+ pd_prefill_forward_durations.append( -+ req.time_stats.get_pd_prefill_forward_duration() -+ ) -+ pd_prefill_transfer_queue_durations.append( -+ req.time_stats.get_pd_prefill_transfer_queue_duration() -+ ) -+ pd_decode_prealloc_durations.append( -+ req.time_stats.get_pd_decode_prealloc_duration() -+ ) -+ pd_decode_transfer_durations.append( -+ req.time_stats.get_pd_decode_transfer_duration() -+ ) -+ pd_decode_forward_durations.append( -+ req.time_stats.get_pd_decode_forward_duration() -+ ) -+ pd_bootstrap_durations.append( -+ req.time_stats.get_pd_bootstrap_duration() -+ ) -+ pd_alloc_waiting_durations.append( -+ req.time_stats.get_pd_alloc_waiting_duration() -+ ) -+ pd_transfer_speeds_gb_s.append( -+ req.time_stats.get_pd_transfer_speed_gb_s() -+ ) -+ pd_transfer_totals_mb.append( -+ req.time_stats.get_pd_transfer_total_mb() -+ ) -+ pd_prefill_retry_counts.append( -+ req.time_stats.get_pd_prefill_retry_count() -+ ) - - if not self.spec_algorithm.is_none(): - spec_verify_ct.append(req.spec_verify_ct) -@@ -1134,7 +1180,7 @@ class SchedulerOutputProcessorMixin: - req.log_time_stats() +@@ -1154,7 +1154,7 @@ class SchedulerOutputProcessorMixin: + dp_ranks = [self.dp_rank] * len(rids) if rids else None # Send to detokenizer - if reqs or is_idle_batch: + if rids or is_idle_batch: - if self.model_config.is_multimodal_gen: - return self.send_to_detokenizer.send_output( -@@ -1149,6 +1195,17 @@ class SchedulerOutputProcessorMixin: - prefill_launch_delay=prefill_launch_delays, - prefill_launch_latency=prefill_launch_latencies, - prefill_finished_ts=prefill_finished_timestamps, -+ pd_prefill_bootstrap_queue_duration=pd_prefill_bootstrap_queue_durations, -+ pd_prefill_forward_duration=pd_prefill_forward_durations, -+ pd_prefill_transfer_queue_duration=pd_prefill_transfer_queue_durations, -+ pd_decode_prealloc_duration=pd_decode_prealloc_durations, -+ pd_decode_transfer_duration=pd_decode_transfer_durations, -+ pd_decode_forward_duration=pd_decode_forward_durations, -+ pd_bootstrap_duration=pd_bootstrap_durations, -+ pd_alloc_waiting_duration=pd_alloc_waiting_durations, -+ pd_transfer_speed_gb_s=pd_transfer_speeds_gb_s, -+ pd_transfer_total_mb=pd_transfer_totals_mb, -+ pd_prefill_retry_count=pd_prefill_retry_counts, - finished_reasons=finished_reasons, - decoded_texts=decoded_texts, - decode_ids=decode_ids_list, -@@ -1198,6 +1255,18 @@ class SchedulerOutputProcessorMixin: - prefill_launch_delays = [] - prefill_launch_latencies = [] - prefill_finished_timestamps = [] -+ profiling_enabled = envs.SLIME_ENABLE_PROFILING.get() -+ pd_prefill_bootstrap_queue_durations = [] if profiling_enabled else None -+ pd_prefill_forward_durations = [] if profiling_enabled else None -+ pd_prefill_transfer_queue_durations = [] if profiling_enabled else None -+ pd_decode_prealloc_durations = [] if profiling_enabled else None -+ pd_decode_transfer_durations = [] if profiling_enabled else None -+ pd_decode_forward_durations = [] if profiling_enabled else None -+ pd_bootstrap_durations = [] if profiling_enabled else None -+ pd_alloc_waiting_durations = [] if profiling_enabled else None -+ pd_transfer_speeds_gb_s = [] if profiling_enabled else None -+ pd_transfer_totals_mb = [] if profiling_enabled else None -+ pd_prefill_retry_counts = [] if profiling_enabled else None - retraction_counts = [] - for req in reqs: - if req.finished(): -@@ -1221,6 +1290,40 @@ class SchedulerOutputProcessorMixin: - prefill_finished_timestamps.append( - req.time_stats.get_prefill_finished_ts() - ) -+ if profiling_enabled: -+ pd_prefill_bootstrap_queue_durations.append( -+ req.time_stats.get_pd_prefill_bootstrap_queue_duration() -+ ) -+ pd_prefill_forward_durations.append( -+ req.time_stats.get_pd_prefill_forward_duration() -+ ) -+ pd_prefill_transfer_queue_durations.append( -+ req.time_stats.get_pd_prefill_transfer_queue_duration() -+ ) -+ pd_decode_prealloc_durations.append( -+ req.time_stats.get_pd_decode_prealloc_duration() -+ ) -+ pd_decode_transfer_durations.append( -+ req.time_stats.get_pd_decode_transfer_duration() -+ ) -+ pd_decode_forward_durations.append( -+ req.time_stats.get_pd_decode_forward_duration() -+ ) -+ pd_bootstrap_durations.append( -+ req.time_stats.get_pd_bootstrap_duration() -+ ) -+ pd_alloc_waiting_durations.append( -+ req.time_stats.get_pd_alloc_waiting_duration() -+ ) -+ pd_transfer_speeds_gb_s.append( -+ req.time_stats.get_pd_transfer_speed_gb_s() -+ ) -+ pd_transfer_totals_mb.append( -+ req.time_stats.get_pd_transfer_total_mb() -+ ) -+ pd_prefill_retry_counts.append( -+ req.time_stats.get_pd_prefill_retry_count() -+ ) - retraction_counts.append(req.retraction_count) - self.send_to_detokenizer.send_output( - BatchEmbeddingOutput( -@@ -1231,6 +1334,17 @@ class SchedulerOutputProcessorMixin: - prefill_launch_delay=prefill_launch_delays, - prefill_launch_latency=prefill_launch_latencies, - prefill_finished_ts=prefill_finished_timestamps, -+ pd_prefill_bootstrap_queue_duration=pd_prefill_bootstrap_queue_durations, -+ pd_prefill_forward_duration=pd_prefill_forward_durations, -+ pd_prefill_transfer_queue_duration=pd_prefill_transfer_queue_durations, -+ pd_decode_prealloc_duration=pd_decode_prealloc_durations, -+ pd_decode_transfer_duration=pd_decode_transfer_durations, -+ pd_decode_forward_duration=pd_decode_forward_durations, -+ pd_bootstrap_duration=pd_bootstrap_durations, -+ pd_alloc_waiting_duration=pd_alloc_waiting_durations, -+ pd_transfer_speed_gb_s=pd_transfer_speeds_gb_s, -+ pd_transfer_total_mb=pd_transfer_totals_mb, -+ pd_prefill_retry_count=pd_prefill_retry_counts, - finished_reasons=finished_reasons, - embeddings=embeddings, - prompt_tokens=prompt_tokens, + BatchTokenIDOutput( + rids=rids, diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py -index 7d08f12b35..afc045da20 100644 +index c02ed7997d..61733c4127 100644 --- a/python/sglang/srt/managers/scheduler_profiler_mixin.py +++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py -@@ -347,7 +347,7 @@ class SchedulerProfilerMixin: +@@ -349,7 +349,7 @@ class SchedulerProfilerMixin: if self.profiler_prefill_ct > self.profiler_target_prefill_ct: if self.profile_in_progress: self.stop_profile(stage=ForwardMode.EXTEND) @@ -1797,7 +1240,7 @@ index 7d08f12b35..afc045da20 100644 if self.profile_in_progress: # force trace flush diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py -index 293a843508..244ea4eb1b 100644 +index abcda67946..a53848b79d 100644 --- a/python/sglang/srt/managers/scheduler_update_weights_mixin.py +++ b/python/sglang/srt/managers/scheduler_update_weights_mixin.py @@ -12,6 +12,7 @@ from sglang.srt.constants import ( @@ -1817,7 +1260,7 @@ index 293a843508..244ea4eb1b 100644 ReleaseMemoryOccupationReqInput, ReleaseMemoryOccupationReqOutput, ResumeMemoryOccupationReqInput, -@@ -114,6 +117,11 @@ class SchedulerUpdateWeightsMixin: +@@ -117,6 +120,11 @@ class SchedulerUpdateWeightsMixin: torch.distributed.barrier(group=self.tp_cpu_group) return UpdateWeightsFromIPCReqOutput(success, message) @@ -1829,7 +1272,7 @@ index 293a843508..244ea4eb1b 100644 def get_weights_by_name(self: Scheduler, recv_req: GetWeightsByNameReqInput): parameter = self.tp_worker.get_weights_by_name(recv_req) return GetWeightsByNameReqOutput(parameter) -@@ -137,6 +145,15 @@ class SchedulerUpdateWeightsMixin: +@@ -140,6 +148,15 @@ class SchedulerUpdateWeightsMixin: self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE) self.flush_cache() @@ -1845,7 +1288,7 @@ index 293a843508..244ea4eb1b 100644 if GPU_MEMORY_TYPE_WEIGHTS in tags: self.stashed_model_static_state = _export_static_state( self.tp_worker.model_runner.model -@@ -177,6 +194,15 @@ class SchedulerUpdateWeightsMixin: +@@ -180,6 +197,15 @@ class SchedulerUpdateWeightsMixin: if GPU_MEMORY_TYPE_KV_CACHE in tags: self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE) @@ -1862,7 +1305,7 @@ index 293a843508..244ea4eb1b 100644 def check_weights(self: Scheduler, recv_req: CheckWeightsReqInput): diff --git a/python/sglang/srt/managers/tokenizer_communicator_mixin.py b/python/sglang/srt/managers/tokenizer_communicator_mixin.py -index f2ffa9909d..6e4d1d460b 100644 +index 544c609401..841658c30e 100644 --- a/python/sglang/srt/managers/tokenizer_communicator_mixin.py +++ b/python/sglang/srt/managers/tokenizer_communicator_mixin.py @@ -59,6 +59,8 @@ from sglang.srt.managers.io_struct import ( @@ -1895,7 +1338,7 @@ index f2ffa9909d..6e4d1d460b 100644 ( GetWeightsByNameReqOutput, self.get_weights_by_name_communicator.handle_recv, -@@ -522,6 +531,17 @@ class TokenizerCommunicatorMixin: +@@ -530,6 +539,17 @@ class TokenizerCommunicatorMixin: return success, message @@ -1911,13 +1354,13 @@ index f2ffa9909d..6e4d1d460b 100644 + return _Communicator.merge_results(results) + async def init_weights_send_group_for_remote_instance( - self, + self: TokenizerManager, obj: InitWeightsSendGroupForRemoteInstanceReqInput, diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 0914a5230b..9114e3e713 100644 +index 81424329a0..2c132be63d 100644 --- a/python/sglang/srt/managers/tokenizer_manager.py +++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -1327,7 +1327,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi +@@ -1383,7 +1383,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): async with self.is_pause_cond: self.is_pause = True if obj.mode != "abort": @@ -1926,7 +1369,7 @@ index 0914a5230b..9114e3e713 100644 else: # we are using the model_update_lock to check if there is still on-going requests. while True: -@@ -1341,7 +1341,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi +@@ -1397,7 +1397,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): async def continue_generation(self, obj: ContinueGenerationReqInput): async with self.is_pause_cond: self.is_pause = False @@ -1935,84 +1378,41 @@ index 0914a5230b..9114e3e713 100644 self.is_pause_cond.notify_all() async def update_weights_from_disk( -@@ -1510,6 +1510,40 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi - self._add_metric_if_present( - recv_obj, "prefill_finished_ts", meta_info, i - ) -+ # PD disaggregation timing -+ self._add_metric_if_present( -+ recv_obj, "pd_prefill_bootstrap_queue_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_prefill_forward_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_prefill_transfer_queue_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_decode_prealloc_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_decode_transfer_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_decode_forward_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_bootstrap_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_alloc_waiting_duration", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_transfer_speed_gb_s", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_transfer_total_mb", meta_info, i -+ ) -+ self._add_metric_if_present( -+ recv_obj, "pd_prefill_retry_count", meta_info, i -+ ) - - if getattr(state.obj, "return_logprob", False): - self.convert_logprob_style( -@@ -1955,19 +1989,17 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi - if custom_labels - else self.metrics_collector.labels - ) +@@ -1965,25 +1965,23 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin): + priority = getattr(state.obj, "priority", None) + if priority is not None: + labels["priority"] = str(priority) - if ( -- state.first_token_time == 0.0 +- not state.ttft_observed - and self.disaggregation_mode != DisaggregationMode.PREFILL - ): -+ if state.first_token_time == 0.0: - state.first_token_time = state.last_time = time.time() - state.first_token_time_perf = time.perf_counter() ++ if not state.ttft_observed: + state.ttft_observed = True state.last_completion_tokens = completion_tokens - self.metrics_collector.observe_time_to_first_token( -- labels, state.first_token_time - state.created_time +- labels, state.time_stats.get_first_token_latency() - ) + if self.disaggregation_mode != DisaggregationMode.PREFILL: + self.metrics_collector.observe_time_to_first_token( -+ labels, state.first_token_time - state.created_time ++ labels, state.time_stats.get_first_token_latency() + ) else: num_new_tokens = completion_tokens - state.last_completion_tokens - if num_new_tokens: + if num_new_tokens > 0: - new_time = time.time() - interval = new_time - state.last_time self.metrics_collector.observe_inter_token_latency( -@@ -1976,7 +2008,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi + labels, + state.time_stats.get_interval(), num_new_tokens, ) - state.last_time = new_time + state.time_stats.set_last_time() - state.last_completion_tokens = completion_tokens + state.last_completion_tokens = completion_tokens if state.finished: retraction_count = ( diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py -index 86b009df4e..16ebd52ae3 100644 +index 7f63610da8..fb56de1583 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -29,6 +29,7 @@ from sglang.srt.managers.io_struct import ( @@ -2023,7 +1423,7 @@ index 86b009df4e..16ebd52ae3 100644 SendWeightsToRemoteInstanceReqInput, UnloadLoRAAdapterReqInput, UpdateWeightFromDiskReqInput, -@@ -168,6 +169,11 @@ class BaseTpWorker(ABC): +@@ -170,6 +171,11 @@ class BaseTpWorker(ABC): success, message = self.model_runner.update_weights_from_ipc(recv_req) return success, message @@ -2035,33 +1435,11 @@ index 86b009df4e..16ebd52ae3 100644 def get_weights_by_name(self, recv_req: GetWeightsByNameReqInput): parameter = self.model_runner.get_weights_by_name( recv_req.name, recv_req.truncate_size -diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py -index fa08bb66a4..22c1c2a127 100644 ---- a/python/sglang/srt/mem_cache/allocator.py -+++ b/python/sglang/srt/mem_cache/allocator.py -@@ -411,7 +411,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): - - self.seen_max_num_extend_tokens_next_power_of_2 = max( - self.seen_max_num_extend_tokens_next_power_of_2, -- min(tl.core.TRITON_MAX_TENSOR_NUMEL, next_power_of_2(extend_num_tokens)), -+ min(65536, next_power_of_2(extend_num_tokens)), - ) - - bs = len(prefix_lens) -@@ -424,7 +424,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): - (extend_num_tokens,), dtype=torch.int64, device=self.device - ) - -- if extend_num_tokens < tl.core.TRITON_MAX_TENSOR_NUMEL: -+ if extend_num_tokens < 65536: - alloc_extend_kernel[(bs,)]( - prefix_lens, - seq_lens, diff --git a/python/sglang/srt/mem_cache/hiradix_cache.py b/python/sglang/srt/mem_cache/hiradix_cache.py -index d7cd472a98..9cf1185cbb 100644 +index 3c1e97daab..e5128e5ee2 100644 --- a/python/sglang/srt/mem_cache/hiradix_cache.py +++ b/python/sglang/srt/mem_cache/hiradix_cache.py -@@ -750,9 +750,8 @@ class HiRadixCache(RadixCache): +@@ -755,9 +755,8 @@ class HiRadixCache(RadixCache): self._update_leaf_status(node) self._update_host_leaf_status(node) if node.parent is None: @@ -2071,9 +1449,9 @@ index d7cd472a98..9cf1185cbb 100644 + # Node belongs to a stale (flushed) tree — stop traversal gracefully. + break node = node.parent - return delta + return DecLockRefResult(delta=delta) -@@ -827,6 +826,7 @@ class HiRadixCache(RadixCache): +@@ -832,6 +831,7 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(node) # update leaf status for the parent because the node is evicted self._update_leaf_status(node.parent) @@ -2081,7 +1459,7 @@ index d7cd472a98..9cf1185cbb 100644 return num_evicted def _evict_regular(self, node: TreeNode): -@@ -1330,6 +1330,7 @@ class HiRadixCache(RadixCache): +@@ -1354,6 +1354,7 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(node) # update parent status as a new leaf is added into device self._update_leaf_status(node.parent) @@ -2089,7 +1467,7 @@ index d7cd472a98..9cf1185cbb 100644 else: self._inc_hit_count(node, chunked) total_prefix_length += prefix_len -@@ -1345,6 +1346,7 @@ class HiRadixCache(RadixCache): +@@ -1369,6 +1370,7 @@ class HiRadixCache(RadixCache): self._update_host_leaf_status(new_node) # update parent status as a new leaf is added into device self._update_leaf_status(new_node.parent) @@ -2098,10 +1476,10 @@ index d7cd472a98..9cf1185cbb 100644 self._inc_hit_count(new_node, chunked) total_prefix_length += prefix_len diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py -index 1d917137c6..669e5c5181 100644 +index e4c158cda9..cf7333235f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py -@@ -1777,9 +1777,12 @@ class NSATokenToKVPool(MLATokenToKVPool): +@@ -1854,9 +1854,12 @@ class NSATokenToKVPool(MLATokenToKVPool): else: assert self.page_size == 64 with ( @@ -2117,19 +1495,19 @@ index 1d917137c6..669e5c5181 100644 ): self.index_k_with_scale_buffer = [ torch.zeros( -@@ -1801,6 +1804,11 @@ class NSATokenToKVPool(MLATokenToKVPool): +@@ -1878,6 +1881,11 @@ class NSATokenToKVPool(MLATokenToKVPool): ) for _ in range(layer_num) ] -+ self.index_k_with_scale_buffer_ptrs = torch.tensor( -+ [x.data_ptr() for x in self.index_k_with_scale_buffer], -+ dtype=torch.uint64, -+ device=self.device, -+ ) ++ self.index_k_with_scale_buffer_ptrs = torch.tensor( ++ [x.data_ptr() for x in self.index_k_with_scale_buffer], ++ dtype=torch.uint64, ++ device=self.device, ++ ) self._finalize_allocation_log(size) def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor: -@@ -1876,6 +1884,50 @@ class NSATokenToKVPool(MLATokenToKVPool): +@@ -1960,6 +1968,50 @@ class NSATokenToKVPool(MLATokenToKVPool): ] return data_ptrs, data_lens, item_lens @@ -2181,10 +1559,10 @@ index 1d917137c6..669e5c5181 100644 kv_size_bytes = super().get_kv_size_bytes() for index_k_cache in self.index_k_with_scale_buffer: diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py -index 42b169728a..8e799196a4 100644 +index 7d16160372..70fbdc702f 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py -@@ -495,7 +495,17 @@ class RadixCache(BasePrefixCache): +@@ -512,7 +512,17 @@ class RadixCache(BasePrefixCache): if self.disable: return @@ -2203,7 +1581,7 @@ index 42b169728a..8e799196a4 100644 kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] -@@ -619,9 +629,8 @@ class RadixCache(BasePrefixCache): +@@ -638,9 +648,8 @@ class RadixCache(BasePrefixCache): node.lock_ref -= 1 self._update_leaf_status(node) if node.parent is None: @@ -2213,253 +1591,70 @@ index 42b169728a..8e799196a4 100644 + # Node belongs to a stale (flushed) tree — stop traversal gracefully. + break node = node.parent - return delta + return DecLockRefResult(delta=delta) -diff --git a/python/sglang/srt/metrics/collector.py b/python/sglang/srt/metrics/collector.py -index 255d41ccc0..f93bedb4dc 100644 ---- a/python/sglang/srt/metrics/collector.py -+++ b/python/sglang/srt/metrics/collector.py -@@ -20,7 +20,10 @@ import time - from dataclasses import dataclass, field - from typing import Any, Dict, List, Optional, Union +diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py +index a59742b943..a7347c15b8 100644 +--- a/python/sglang/srt/model_executor/model_runner.py ++++ b/python/sglang/srt/model_executor/model_runner.py +@@ -406,7 +406,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): + self.forward_stream = torch.get_device_module(self.device).Stream() --from sglang.srt.disaggregation.utils import DisaggregationMode -+from sglang.srt.disaggregation.utils import ( -+ DisaggregationMode, -+ is_slime_profiling_enabled, -+) - from sglang.srt.environ import envs - from sglang.srt.metrics.utils import exponential_buckets, generate_buckets - from sglang.srt.model_executor.forward_batch_info import ForwardMode -@@ -77,6 +80,17 @@ class TimeStats: - # Number of prefill retries for this request - prefill_retry_count: int = 0 + # CPU offload +- set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank)) ++ # For draft worker (e.g., MTP), do not set offloader to avoid overriding ++ # the main model's offloader. Draft worker uses NoopOffloader instead. ++ if not is_draft_worker: ++ set_offloader( ++ create_offloader_from_server_args(server_args, dp_rank=dp_rank) ++ ) -+ # Prefill-side durations forwarded via metadata transfer from P to D instance. -+ # Set on the decode instance after KV cache transfer completes. -+ fwd_prefill_bootstrap_queue_duration: Optional[float] = None -+ fwd_prefill_forward_duration: Optional[float] = None -+ fwd_prefill_transfer_queue_duration: Optional[float] = None -+ fwd_bootstrap_duration: Optional[float] = None -+ fwd_alloc_waiting_duration: Optional[float] = None -+ fwd_transfer_speed_gb_s: Optional[float] = None -+ fwd_transfer_total_mb: Optional[float] = None -+ fwd_prefill_retry_count: Optional[int] = None -+ - # Timestamp when prefill phase finishes, obtained from `time.time()`. - # Note that this differs from the other `_time` fields tracked by the - # `TimeStats` class, which are obtained from `time.perf_counter()`. -@@ -102,6 +116,148 @@ class TimeStats: - return self.prefill_finished_ts - return None + self._weight_checker = WeightChecker(model_runner=self) + +@@ -646,7 +651,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): + ) + + # Init routed experts capturer +- self.init_routed_experts_capturer() ++ if not self.is_draft_worker: ++ self.init_routed_experts_capturer() + + if self.device == "cuda" or self.device == "musa": + self.init_cublas() +@@ -2767,11 +2773,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): + output.expert_distribution_metrics = recorder_outputs.get("metrics") -+ # --- PD disaggregation timing getters --- -+ -+ def get_pd_prefill_bootstrap_queue_duration(self) -> Optional[float]: -+ """P instance: time spent in bootstrap queue before entering the wait queue.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_prefill_bootstrap_queue_duration is not None: -+ return self.fwd_prefill_bootstrap_queue_duration -+ if ( -+ self.disagg_mode == DisaggregationMode.PREFILL -+ and self.prefill_bootstrap_queue_entry_time > 0.0 -+ and self.wait_queue_entry_time > 0.0 -+ ): -+ return self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time -+ return None -+ -+ def get_pd_prefill_forward_duration(self) -> Optional[float]: -+ """P instance: time for the actual prefill forward computation.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_prefill_forward_duration is not None: -+ return self.fwd_prefill_forward_duration -+ if ( -+ self.disagg_mode == DisaggregationMode.PREFILL -+ and self.forward_entry_time > 0.0 -+ and self.completion_time > 0.0 -+ ): -+ return self.completion_time - self.forward_entry_time -+ return None -+ -+ def get_pd_prefill_transfer_queue_duration(self) -> Optional[float]: -+ """P instance: time spent in the transfer queue (KV cache send).""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_prefill_transfer_queue_duration is not None: -+ return self.fwd_prefill_transfer_queue_duration -+ if ( -+ self.disagg_mode == DisaggregationMode.PREFILL -+ and self.prefill_transfer_queue_entry_time > 0.0 -+ and self.completion_time > 0.0 -+ ): -+ return self.completion_time - self.prefill_transfer_queue_entry_time -+ return None -+ -+ def get_pd_decode_prealloc_duration(self) -> Optional[float]: -+ """D instance: time spent in the pre-alloc queue (waiting for KV cache slot allocation).""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if ( -+ self.disagg_mode == DisaggregationMode.DECODE -+ and self.decode_prealloc_queue_entry_time > 0.0 -+ and self.decode_transfer_queue_entry_time > 0.0 -+ ): -+ return ( -+ self.decode_transfer_queue_entry_time -+ - self.decode_prealloc_queue_entry_time -+ ) -+ return None -+ -+ def get_pd_decode_transfer_duration(self) -> Optional[float]: -+ """D instance: time spent waiting for KV cache transfer to complete.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if ( -+ self.disagg_mode == DisaggregationMode.DECODE -+ and self.decode_transfer_queue_entry_time > 0.0 -+ and self.wait_queue_entry_time > 0.0 -+ ): -+ return self.wait_queue_entry_time - self.decode_transfer_queue_entry_time -+ return None -+ -+ def get_pd_decode_forward_duration(self) -> Optional[float]: -+ """D instance: time for the actual decode forward computation.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if ( -+ self.disagg_mode == DisaggregationMode.DECODE -+ and self.forward_entry_time > 0.0 -+ and self.completion_time > 0.0 -+ ): -+ return self.completion_time - self.forward_entry_time -+ return None -+ -+ def get_pd_bootstrap_duration(self) -> Optional[float]: -+ """Bootstrap handshake duration (both P and D instances).""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_bootstrap_duration is not None: -+ return self.fwd_bootstrap_duration -+ if ( -+ self.disagg_mode != DisaggregationMode.NULL -+ and self.bootstrap_duration > 0.0 -+ ): -+ return self.bootstrap_duration -+ return None -+ -+ def get_pd_alloc_waiting_duration(self) -> Optional[float]: -+ """KV cache allocation waiting duration (both P and D instances).""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_alloc_waiting_duration is not None: -+ return self.fwd_alloc_waiting_duration -+ if ( -+ self.disagg_mode != DisaggregationMode.NULL -+ and self.alloc_waiting_duration > 0.0 -+ ): -+ return self.alloc_waiting_duration -+ return None -+ -+ def get_pd_transfer_speed_gb_s(self) -> Optional[float]: -+ """KV cache transfer speed in GB/s.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_transfer_speed_gb_s is not None: -+ return self.fwd_transfer_speed_gb_s -+ if ( -+ self.disagg_mode != DisaggregationMode.NULL -+ and self.transfer_speed_gb_s > 0.0 -+ ): -+ return self.transfer_speed_gb_s -+ return None -+ -+ def get_pd_transfer_total_mb(self) -> Optional[float]: -+ """Total KV cache transferred in MB.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_transfer_total_mb is not None: -+ return self.fwd_transfer_total_mb -+ if self.disagg_mode != DisaggregationMode.NULL and self.transfer_total_mb > 0.0: -+ return self.transfer_total_mb -+ return None -+ -+ def get_pd_prefill_retry_count(self) -> Optional[int]: -+ """Number of prefill retries for this request.""" -+ if not is_slime_profiling_enabled(): -+ return None -+ if self.fwd_prefill_retry_count is not None: -+ return self.fwd_prefill_retry_count -+ if self.disagg_mode == DisaggregationMode.PREFILL: -+ return self.prefill_retry_count -+ return None -+ - def convert_to_duration(self) -> str: - if self.disagg_mode == DisaggregationMode.NULL: - queue_duration = self.forward_entry_time - self.wait_queue_entry_time -diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py -index 275775a73d..e4e2fdc398 100644 ---- a/python/sglang/srt/model_executor/model_runner.py -+++ b/python/sglang/srt/model_executor/model_runner.py -@@ -395,7 +395,12 @@ class ModelRunner(ModelRunnerKVCacheMixin): - self.forward_stream = torch.get_device_module(self.device).Stream() - - # CPU offload -- set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank)) -+ # For draft worker (e.g., MTP), do not set offloader to avoid overriding -+ # the main model's offloader. Draft worker uses NoopOffloader instead. -+ if not is_draft_worker: -+ set_offloader( -+ create_offloader_from_server_args(server_args, dp_rank=dp_rank) -+ ) - - self._weight_checker = WeightChecker(model_runner=self) - -@@ -600,7 +605,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): - ) - - # Init routed experts capturer -- self.init_routed_experts_capturer() -+ if not self.is_draft_worker: -+ self.init_routed_experts_capturer() - - if self.device == "cuda" or self.device == "musa": - self.init_cublas() -@@ -2429,11 +2435,19 @@ class ModelRunner(ModelRunnerKVCacheMixin): - output.expert_distribution_metrics = recorder_outputs.get("metrics") - - # Copy cached routing experts' buffers back to CPU cache -- get_global_experts_capturer().on_forward_end( -- forward_batch=forward_batch, -- can_run_graph=output.can_run_graph, -- cuda_graph_batch=getattr(self.graph_runner, "bs", None), -- ) -+ if not self.is_draft_worker: -+ # In speculative decoding, num_tokens_per_bs > 1, so we need to pass -+ # the actual number of tokens per dp rank in cuda graph, not batch size. -+ cuda_graph_num_tokens = None -+ if getattr(self.graph_runner, "bs", None): -+ cuda_graph_num_tokens = ( -+ self.graph_runner.bs * self.graph_runner.num_tokens_per_bs -+ ) -+ get_global_experts_capturer().on_forward_end( -+ forward_batch=forward_batch, -+ can_run_graph=output.can_run_graph, -+ cuda_graph_batch=cuda_graph_num_tokens, -+ ) - - if self.eplb_manager is not None: - self.eplb_manager.on_forward_pass_end() -@@ -2664,6 +2678,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): - device=self.device, - ) - -+ def post_process_weights(self, recv_req): -+ """ -+ Execute post-processing logic for model weights, such as Marlin quantization format conversion. -+ """ -+ from sglang.srt.model_loader.loader import device_loading_context + # Copy cached routing experts' buffers back to CPU cache +- get_global_experts_capturer().on_forward_end( +- forward_batch=forward_batch, +- can_run_graph=output.can_run_graph, +- cuda_graph_batch=getattr(self.graph_runner, "bs", None), +- ) ++ if not self.is_draft_worker: ++ # In speculative decoding, num_tokens_per_bs > 1, so we need to pass ++ # the actual number of tokens per dp rank in cuda graph, not batch size. ++ cuda_graph_num_tokens = None ++ if getattr(self.graph_runner, "bs", None): ++ cuda_graph_num_tokens = ( ++ self.graph_runner.bs * self.graph_runner.num_tokens_per_bs ++ ) ++ get_global_experts_capturer().on_forward_end( ++ forward_batch=forward_batch, ++ can_run_graph=output.can_run_graph, ++ cuda_graph_batch=cuda_graph_num_tokens, ++ ) + + if self.eplb_manager is not None: + self.eplb_manager.on_forward_pass_end() +@@ -3021,6 +3035,42 @@ class ModelRunner(ModelRunnerKVCacheMixin): + device=self.device, + ) + ++ def post_process_weights(self, recv_req): ++ """ ++ Execute post-processing logic for model weights, such as Marlin quantization format conversion. ++ """ ++ from sglang.srt.model_loader.loader import device_loading_context + + target_device = torch.device("cuda", torch.cuda.current_device()) + @@ -2494,320 +1689,8 @@ index 275775a73d..e4e2fdc398 100644 def _model_load_weights_direct(model, named_tensors: List[Tuple[str, torch.Tensor]]): params_dict = dict(model.named_parameters()) -diff --git a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py -index cc673a9cac..06c430d2c4 100644 ---- a/python/sglang/srt/models/deepseek_common/attention_backend_handler.py -+++ b/python/sglang/srt/models/deepseek_common/attention_backend_handler.py -@@ -1,4 +1,5 @@ - from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph -+from sglang.srt.layers.attention.hybrid_attn_backend import HybridAttnBackend - from sglang.srt.layers.attention.tbo_backend import TboAttnBackend - from sglang.srt.models.deepseek_common.attention_forward_methods.forward_methods import ( - AttnForwardMethod, -@@ -150,6 +151,8 @@ def handle_attention_nsa(attn, forward_batch): - backend = forward_batch.attn_backend - if isinstance(backend, TboAttnBackend): # if enable tbo, get primary backend - backend = backend.primary -+ if isinstance(backend, HybridAttnBackend): -+ backend = backend._select_backend(forward_batch.forward_mode) - if hasattr(backend, "use_mha") and backend.use_mha: - return AttnForwardMethod.MHA_ONE_SHOT - return AttnForwardMethod.MLA -diff --git a/python/sglang/srt/models/deepseek_nextn.py b/python/sglang/srt/models/deepseek_nextn.py -index cb13a7c676..d9669ce086 100644 ---- a/python/sglang/srt/models/deepseek_nextn.py -+++ b/python/sglang/srt/models/deepseek_nextn.py -@@ -29,6 +29,7 @@ from sglang.srt.layers.attention.nsa.utils import ( - can_cp_split, - cp_all_gather_rerange_output, - cp_split_and_rebuild_data, -+ cp_split_and_rebuild_position, - is_nsa_enable_prefill_cp, - nsa_use_prefill_cp, - prepare_input_dp_with_cp_dsa, -@@ -160,15 +161,17 @@ class DeepseekModelNextN(nn.Module): - - if nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp): - hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states) -+ positions = cp_split_and_rebuild_position(forward_batch, positions) - residual = None - with get_global_expert_distribution_recorder().disable_this_region(): -- hidden_states, residual = self.decoder( -+ hidden_states, residual, *rest = self.decoder( - positions, - hidden_states, - forward_batch, - residual, - zero_allocator, - ) -+ topk_indices = rest[0] if rest else None - - if not forward_batch.forward_mode.is_idle(): - if residual is not None: -diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py -index 1583dd7880..a35c00f96c 100644 ---- a/python/sglang/srt/models/deepseek_v2.py -+++ b/python/sglang/srt/models/deepseek_v2.py -@@ -1085,6 +1085,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - prefix: str = "", - alt_stream: Optional[torch.cuda.Stream] = None, - skip_rope: bool = False, -+ is_nextn: bool = False, - ) -> None: - super().__init__() - self.layer_id = layer_id -@@ -1154,6 +1155,8 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - prefix=add_prefix("kv_a_proj_with_mqa", prefix), - ) - -+ self.skip_topk = False -+ self.next_skip_topk = False - if self.use_nsa: - is_neox_style = not getattr(config, "indexer_rope_interleave", False) - self.indexer = Indexer( -@@ -1174,6 +1177,31 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - layer_id=layer_id, - alt_stream=alt_stream, - ) -+ if not is_nextn: -+ self.index_topk_freq = getattr(config, "index_topk_freq", 1) -+ self.index_topk_pattern = getattr(config, "index_topk_pattern", None) -+ self.index_skip_topk_offset = getattr( -+ config, "index_skip_topk_offset", 2 -+ ) -+ if self.index_topk_pattern is None: -+ self.skip_topk = ( -+ max(layer_id - self.index_skip_topk_offset + 1, 0) -+ % self.index_topk_freq -+ != 0 -+ ) -+ self.next_skip_topk = ( -+ max(layer_id - self.index_skip_topk_offset + 2, 0) -+ % self.index_topk_freq -+ != 0 -+ ) -+ else: -+ self.skip_topk = self.index_topk_pattern[layer_id] == "S" -+ if layer_id < len(self.index_topk_pattern) - 1: -+ self.next_skip_topk = ( -+ self.index_topk_pattern[layer_id + 1] == "S" -+ ) -+ else: -+ self.next_skip_topk = False - - self.kv_b_proj = ColumnParallelLinear( - self.kv_lora_rank, -@@ -1362,6 +1390,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - llama_4_scaling: Optional[torch.Tensor] = None, -+ prev_topk_indices: Optional[torch.Tensor] = None, - ): - s = self.forward_prepare( - positions=positions, -@@ -1369,6 +1398,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - forward_batch=forward_batch, - zero_allocator=zero_allocator, - llama_4_scaling=llama_4_scaling, -+ prev_topk_indices=prev_topk_indices, - ) - return self.forward_core(s) - -@@ -1379,6 +1409,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - llama_4_scaling: Optional[torch.Tensor] = None, -+ prev_topk_indices: Optional[torch.Tensor] = None, - ): - if self.attn_mha.kv_b_proj is None: - self.attn_mha.kv_b_proj = self.kv_b_proj -@@ -1418,7 +1449,12 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - ) - elif attn_forward_method == AttnForwardMethod.MLA: - inner_state = self.forward_absorb_prepare( -- positions, hidden_states, forward_batch, zero_allocator, llama_4_scaling -+ positions, -+ hidden_states, -+ forward_batch, -+ zero_allocator, -+ llama_4_scaling, -+ prev_topk_indices, - ) - elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE: - inner_state = self.forward_absorb_fused_mla_rope_prepare( -@@ -1529,6 +1565,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - forward_batch: ForwardBatch, - zero_allocator: BumpAllocator, - llama_4_scaling: Optional[torch.Tensor] = None, -+ prev_topk_indices: Optional[torch.Tensor] = None, - ): - from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode - -@@ -1620,18 +1657,7 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - q = self.q_b_proj(q)[0].view( - -1, self.num_local_heads, self.qk_head_dim - ) -- topk_indices = self.indexer( -- x=hidden_states, -- q_lora=q_lora, -- positions=positions, -- forward_batch=forward_batch, -- layer_id=self.layer_id, -- ) -- current_stream.wait_stream(self.alt_stream) -- else: -- k_nope = k_nope.unsqueeze(1) -- q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) -- if q_lora is not None: -+ if not self.skip_topk: - topk_indices = self.indexer( - x=hidden_states, - q_lora=q_lora, -@@ -1639,6 +1665,23 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - forward_batch=forward_batch, - layer_id=self.layer_id, - ) -+ else: -+ topk_indices = prev_topk_indices -+ current_stream.wait_stream(self.alt_stream) -+ else: -+ k_nope = k_nope.unsqueeze(1) -+ q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim) -+ if q_lora is not None: -+ if not self.skip_topk: -+ topk_indices = self.indexer( -+ x=hidden_states, -+ q_lora=q_lora, -+ positions=positions, -+ forward_batch=forward_batch, -+ layer_id=self.layer_id, -+ ) -+ else: -+ topk_indices = prev_topk_indices - else: - q = self.q_proj(hidden_states)[0].view( - -1, self.num_local_heads, self.qk_head_dim -@@ -1929,8 +1972,10 @@ class DeepseekV2AttentionMLA(nn.Module, DeepseekMHAForwardMixin): - ).transpose(0, 1), - ) - output, _ = self.o_proj(attn_bmm_output) -- -- return output -+ if not self.next_skip_topk: -+ return output, None -+ else: -+ return output, topk_indices - - def forward_absorb_fused_mla_rope_prepare( - self, -@@ -2275,6 +2320,7 @@ class DeepseekV2DecoderLayer(nn.Module): - reduce_results=False, - prefix=add_prefix("self_attn", prefix), - alt_stream=alt_stream, -+ is_nextn=is_nextn, - ) - - self.is_layer_sparse = self._is_layer_sparse(layer_id, is_nextn=is_nextn) -@@ -2357,6 +2403,7 @@ class DeepseekV2DecoderLayer(nn.Module): - zero_allocator: BumpAllocator, - gemm_output_zero_allocator: BumpAllocator = None, - llama_4_scaling: Optional[torch.Tensor] = None, -+ prev_topk_indices: Optional[torch.Tensor] = None, - ) -> torch.Tensor: - quant_format = ( - "mxfp4" -@@ -2398,7 +2445,12 @@ class DeepseekV2DecoderLayer(nn.Module): - forward_batch=forward_batch, - zero_allocator=zero_allocator, - llama_4_scaling=llama_4_scaling, -+ prev_topk_indices=prev_topk_indices, - ) -+ if isinstance(hidden_states, tuple): -+ hidden_states, topk_indices = hidden_states -+ else: -+ topk_indices = None - - hidden_states, residual = self.layer_communicator.prepare_mlp( - hidden_states, residual, forward_batch -@@ -2434,7 +2486,7 @@ class DeepseekV2DecoderLayer(nn.Module): - hidden_states, residual, forward_batch - ) - -- return hidden_states, residual -+ return hidden_states, residual, topk_indices - - def op_comm_prepare_attn( - self, -@@ -2710,6 +2762,7 @@ class DeepseekV2Model(nn.Module): - elif self.first_k_dense_replace < normal_start_layer: - normal_end_layer = normal_start_layer = 0 - aux_hidden_states = [] -+ topk_indices = None - for i in range(normal_start_layer, normal_end_layer): - # NOTE: torch dynamo does not support graph break in context manager - ctx = ( -@@ -2727,7 +2780,7 @@ class DeepseekV2Model(nn.Module): - else: - aux_hidden_states.append(hidden_states + residual) - layer = self.layers[i] -- hidden_states, residual = layer( -+ hidden_states, residual, *rest = layer( - positions, - hidden_states, - forward_batch, -@@ -2735,7 +2788,9 @@ class DeepseekV2Model(nn.Module): - zero_allocator, - gemm_output_zero_allocator, - llama_4_scaling, -+ prev_topk_indices=topk_indices, - ) -+ topk_indices = rest[0] if rest else None - - if normal_end_layer != self.end_layer: - hidden_states, residual = model_forward_maybe_tbo( -diff --git a/python/sglang/srt/models/glm4_moe.py b/python/sglang/srt/models/glm4_moe.py -index db8c1c7ce7..53ffadf6d0 100644 ---- a/python/sglang/srt/models/glm4_moe.py -+++ b/python/sglang/srt/models/glm4_moe.py -@@ -678,8 +678,13 @@ class Glm4MoeDecoderLayer(nn.Module): - nn.Module.__init__(self) - self.hidden_size = config.hidden_size - self.config = config -- rope_theta = getattr(config, "rope_theta", 10000) -- rope_scaling = getattr(config, "rope_scaling", None) -+ # rope_theta may be stored in rope_parameters dict (e.g. GLM-4.6V) -+ _rope_params = getattr(config, "rope_parameters", None) -+ if isinstance(_rope_params, dict) and "rope_theta" in _rope_params: -+ rope_theta = _rope_params["rope_theta"] -+ else: -+ rope_theta = getattr(config, "rope_theta", 10000) -+ rope_scaling = getattr(config, "rope_scaling", None) or _rope_params - partial_rotary_factor = getattr( - getattr(config, "rope_parameters", None), "partial_rotary_factor", None - ) or getattr(config, "partial_rotary_factor", 0.5) -@@ -773,6 +778,7 @@ class Glm4MoeDecoderLayer(nn.Module): - hidden_states: torch.Tensor, - forward_batch: ForwardBatch, - residual: Optional[torch.Tensor], -+ **kwargs, - ) -> torch.Tensor: - - hidden_states, residual = self.layer_communicator.prepare_attn( -diff --git a/python/sglang/srt/models/glm4_moe_nextn.py b/python/sglang/srt/models/glm4_moe_nextn.py -index 1f6e753646..546cce4ab5 100644 ---- a/python/sglang/srt/models/glm4_moe_nextn.py -+++ b/python/sglang/srt/models/glm4_moe_nextn.py -@@ -103,7 +103,7 @@ class Glm4MoeModelNextN(nn.Module): - - residual = None - with get_global_expert_distribution_recorder().disable_this_region(): -- hidden_states, residual = self.decoder( -+ hidden_states, residual, *rest = self.decoder( - positions, hidden_states, forward_batch, residual - ) - diff --git a/python/sglang/srt/models/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py -index 324de18b49..fc72faa031 100644 +index 2f0074924d..1f991932c6 100644 --- a/python/sglang/srt/models/glm4v_moe.py +++ b/python/sglang/srt/models/glm4v_moe.py @@ -52,11 +52,31 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration): @@ -2911,87 +1794,24 @@ index 324de18b49..fc72faa031 100644 if name not in params_dict: continue -diff --git a/python/sglang/srt/models/qwen3_5.py b/python/sglang/srt/models/qwen3_5.py -index f01225487b..1dad8bb8e5 100644 ---- a/python/sglang/srt/models/qwen3_5.py -+++ b/python/sglang/srt/models/qwen3_5.py -@@ -372,6 +372,7 @@ class Qwen3_5LinearDecoderLayer(nn.Module): - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, -+ is_last_layer=(layer_id == config.num_hidden_layers - 1), - ) - - def forward( -@@ -400,11 +401,24 @@ class Qwen3_5LinearDecoderLayer(nn.Module): - use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( - forward_batch - ) -- hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) - -- hidden_states, residual = self.layer_communicator.postprocess_layer( -- hidden_states, residual, forward_batch -+ should_allreduce_fusion = ( -+ self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( -+ forward_batch -+ ) - ) -+ if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): -+ hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) -+ else: -+ hidden_states = self.mlp( -+ hidden_states, should_allreduce_fusion, use_reduce_scatter -+ ) -+ if should_allreduce_fusion: -+ hidden_states._sglang_needs_allreduce_fusion = True -+ else: -+ hidden_states, residual = self.layer_communicator.postprocess_layer( -+ hidden_states, residual, forward_batch -+ ) - - return hidden_states, residual - -@@ -549,6 +563,7 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): - input_layernorm=self.input_layernorm, - post_attention_layernorm=self.post_attention_layernorm, - allow_reduce_scatter=True, -+ is_last_layer=(layer_id == config.num_hidden_layers - 1), - ) - - self.alt_stream = alt_stream -@@ -633,11 +648,24 @@ class Qwen3_5AttentionDecoderLayer(nn.Module): - use_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( - forward_batch - ) -- hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) - -- hidden_states, residual = self.layer_communicator.postprocess_layer( -- hidden_states, residual, forward_batch -+ should_allreduce_fusion = ( -+ self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( -+ forward_batch -+ ) - ) -+ if isinstance(self.mlp, Qwen2MoeSparseMoeBlock): -+ hidden_states = self.mlp(hidden_states, forward_batch, use_reduce_scatter) -+ else: -+ hidden_states = self.mlp( -+ hidden_states, should_allreduce_fusion, use_reduce_scatter -+ ) -+ if should_allreduce_fusion: -+ hidden_states._sglang_needs_allreduce_fusion = True -+ else: -+ hidden_states, residual = self.layer_communicator.postprocess_layer( -+ hidden_states, residual, forward_batch -+ ) +diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py +index 912891b6a7..fd67a7b580 100644 +--- a/python/sglang/srt/models/qwen3_moe.py ++++ b/python/sglang/srt/models/qwen3_moe.py +@@ -325,7 +325,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module): + topk_output = self.topk(hidden_states, router_logits) + final_hidden_states = self.experts(hidden_states, topk_output) - return hidden_states, residual +- if self.ep_size > 1 and not should_allreduce_fusion: ++ if self.ep_size > 1 and not should_allreduce_fusion and not use_reduce_scatter: + final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) + if ( diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py -index d641826e33..3abc39ef32 100644 +index 7746b24459..57b65fe06f 100644 --- a/python/sglang/srt/models/qwen3_vl.py +++ b/python/sglang/srt/models/qwen3_vl.py -@@ -711,14 +711,19 @@ class Qwen3LLMModel(Qwen3Model): +@@ -1005,14 +1005,19 @@ class Qwen3LLMModel(Qwen3Model): hidden_states + residual if residual is not None else hidden_states ) @@ -3016,20 +1836,25 @@ index d641826e33..3abc39ef32 100644 positions, hidden_states, diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py -index 33cce6fe25..0970c4550d 100644 +index a44f14b6ca..6d6c65ea49 100644 --- a/python/sglang/srt/multimodal/processors/glm4v.py +++ b/python/sglang/srt/multimodal/processors/glm4v.py -@@ -1,6 +1,9 @@ +@@ -1,7 +1,13 @@ from typing import List, Union +import torch + from sglang.srt.layers.rotary_embedding import MRotaryEmbedding -+from sglang.srt.managers.schedule_batch import Modality, MultimodalDataItem +-from sglang.srt.managers.schedule_batch import MultimodalProcessorOutput ++from sglang.srt.managers.schedule_batch import ( ++ Modality, ++ MultimodalDataItem, ++ MultimodalProcessorOutput, ++) from sglang.srt.models.glm4v import Glm4vForConditionalGeneration from sglang.srt.models.glm4v_moe import Glm4vMoeForConditionalGeneration from sglang.srt.multimodal.processors.base_processor import ( -@@ -45,6 +48,8 @@ class Glm4vImageProcessor(SGLangBaseProcessor): +@@ -46,6 +52,8 @@ class Glm4vImageProcessor(SGLangBaseProcessor): self.IMAGE_END_TOKEN_ID = hf_config.image_end_token_id self.VIDEO_START_TOKEN_ID = hf_config.video_start_token_id self.VIDEO_END_TOKEN_ID = hf_config.video_end_token_id @@ -3038,48 +1863,51 @@ index 33cce6fe25..0970c4550d 100644 # Vision config self.IMAGE_FACTOR = 28 -@@ -59,6 +64,36 @@ class Glm4vImageProcessor(SGLangBaseProcessor): +@@ -60,6 +68,39 @@ class Glm4vImageProcessor(SGLangBaseProcessor): video_token_id=self.IM_TOKEN_ID, ).build(_processor) + def get_mm_data(self, prompt, embeddings, img_grid_thw): -+ input_ids, offsets = self.build_input_ids(prompt, img_grid_thw) ++ input_ids, offsets, _ = self.build_input_ids(prompt, img_grid_thw=img_grid_thw) ++ image_embeddings = ( ++ embeddings.get(Modality.IMAGE, embeddings) ++ if isinstance(embeddings, dict) ++ else embeddings ++ ) + mm_items = [ + MultimodalDataItem( + modality=Modality.IMAGE, + offsets=offsets, -+ precomputed_embeddings=embeddings, ++ precomputed_embeddings=image_embeddings, + ) + ] + -+ input_ids_tensor = torch.tensor(input_ids) + mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_glm4v( -+ input_ids=input_ids_tensor.unsqueeze(0), ++ input_ids=torch.tensor(input_ids, dtype=torch.long).unsqueeze(0), + hf_config=self.hf_config, + image_grid_thw=img_grid_thw, + video_grid_thw=None, + attention_mask=None, + ) -+ mrope_positions = mrope_positions.squeeze(1) -+ -+ return { -+ "input_ids": input_ids, -+ "mm_items": mm_items, -+ "im_start_id": self.IM_START_TOKEN_ID, -+ "im_end_id": self.IM_END_TOKEN_ID, -+ "im_token_id": self.IM_TOKEN_ID, -+ "mrope_positions": mrope_positions, -+ "mrope_position_delta": mrope_position_delta, -+ } + - async def process_mm_data_async( - self, - image_data: List[Union[str, bytes]], ++ return MultimodalProcessorOutput( ++ input_ids=input_ids, ++ mm_items=mm_items, ++ im_start_id=self.IM_START_TOKEN_ID, ++ im_end_id=self.IM_END_TOKEN_ID, ++ im_token_id=self.IM_TOKEN_ID, ++ mrope_positions=mrope_positions.squeeze(1), ++ mrope_position_delta=mrope_position_delta, ++ ) ++ + def compute_mrope_positions(self, input_ids, mm_items): + image_grid_thw = None + video_grid_thw = None diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py -index 4395654e4e..f9b5ea4abb 100644 +index 3f102567d0..6fb3899021 100644 --- a/python/sglang/srt/multimodal/processors/qwen_vl.py +++ b/python/sglang/srt/multimodal/processors/qwen_vl.py -@@ -317,7 +317,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): +@@ -499,7 +499,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor): **kwargs, ): entry_time = time.perf_counter() @@ -3088,19 +1916,306 @@ index 4395654e4e..f9b5ea4abb 100644 prompt=input_text, image_data=image_data, video_data=request_obj.video_data, +diff --git a/python/sglang/srt/observability/req_time_stats.py b/python/sglang/srt/observability/req_time_stats.py +index 8caf21c320..51d1edc584 100644 +--- a/python/sglang/srt/observability/req_time_stats.py ++++ b/python/sglang/srt/observability/req_time_stats.py +@@ -21,7 +21,10 @@ import uuid + from dataclasses import dataclass, field + from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union + +-from sglang.srt.disaggregation.utils import DisaggregationMode ++from sglang.srt.disaggregation.utils import ( ++ DisaggregationMode, ++ is_slime_profiling_enabled, ++) + from sglang.srt.model_executor.forward_batch_info import ForwardMode + from sglang.srt.observability.metrics_collector import ( + SchedulerMetricsCollector, +@@ -553,6 +556,14 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + transfer_total_mb: float = 0.0 + # Number of prefill retries for this request + prefill_retry_count: int = 0 ++ fwd_prefill_bootstrap_queue_duration: Optional[float] = None ++ fwd_prefill_forward_duration: Optional[float] = None ++ fwd_prefill_transfer_queue_duration: Optional[float] = None ++ fwd_bootstrap_duration: Optional[float] = None ++ fwd_alloc_waiting_duration: Optional[float] = None ++ fwd_transfer_speed_gb_s: Optional[float] = None ++ fwd_transfer_total_mb: Optional[float] = None ++ fwd_prefill_retry_count: Optional[int] = None + + def __getstate__(self) -> object: + # send to detokenizer/tokenizer +@@ -560,11 +571,33 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + return {} + + state = { ++ "disagg_mode": self.disagg_mode, + "wait_queue_entry_time": self.wait_queue_entry_time, + "forward_entry_time": self.forward_entry_time, + "prefill_run_batch_start_time": self.prefill_run_batch_start_time, + "prefill_run_batch_end_time": self.prefill_run_batch_end_time, + "prefill_finished_time": self.prefill_finished_time, ++ "completion_time": self.completion_time, ++ "prefill_bootstrap_queue_entry_time": ( ++ self.prefill_bootstrap_queue_entry_time ++ ), ++ "prefill_transfer_queue_entry_time": self.prefill_transfer_queue_entry_time, ++ "decode_prealloc_queue_entry_time": self.decode_prealloc_queue_entry_time, ++ "decode_transfer_queue_entry_time": self.decode_transfer_queue_entry_time, ++ "bootstrap_done_time": self.bootstrap_done_time, ++ "transfer_speed_gb_s": self.transfer_speed_gb_s, ++ "transfer_total_mb": self.transfer_total_mb, ++ "prefill_retry_count": self.prefill_retry_count, ++ "fwd_prefill_bootstrap_queue_duration": ( ++ self.fwd_prefill_bootstrap_queue_duration ++ ), ++ "fwd_prefill_forward_duration": self.fwd_prefill_forward_duration, ++ "fwd_prefill_transfer_queue_duration": self.fwd_prefill_transfer_queue_duration, ++ "fwd_bootstrap_duration": self.fwd_bootstrap_duration, ++ "fwd_alloc_waiting_duration": self.fwd_alloc_waiting_duration, ++ "fwd_transfer_speed_gb_s": self.fwd_transfer_speed_gb_s, ++ "fwd_transfer_total_mb": self.fwd_transfer_total_mb, ++ "fwd_prefill_retry_count": self.fwd_prefill_retry_count, + "diff_realtime_monotonic": global_diff_realtime_monotonic, + } + return state +@@ -916,6 +949,149 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + return self.prefill_run_batch_end_time - self.prefill_run_batch_start_time + return None + ++ def get_pd_prefill_bootstrap_queue_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_bootstrap_queue_duration is not None: ++ return self.fwd_prefill_bootstrap_queue_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.prefill_bootstrap_queue_entry_time > 0.0 ++ and self.wait_queue_entry_time > 0.0 ++ ): ++ return self.wait_queue_entry_time - self.prefill_bootstrap_queue_entry_time ++ return None ++ ++ def get_pd_prefill_forward_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_forward_duration is not None: ++ return self.fwd_prefill_forward_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.forward_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.forward_entry_time ++ return None ++ ++ def get_pd_prefill_transfer_queue_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_transfer_queue_duration is not None: ++ return self.fwd_prefill_transfer_queue_duration ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.prefill_transfer_queue_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.prefill_transfer_queue_entry_time ++ return None ++ ++ def get_pd_decode_prealloc_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_prealloc_queue_entry_time > 0.0 ++ and self.decode_transfer_queue_entry_time > 0.0 ++ ): ++ return ( ++ self.decode_transfer_queue_entry_time ++ - self.decode_prealloc_queue_entry_time ++ ) ++ return None ++ ++ def get_pd_decode_transfer_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_transfer_queue_entry_time > 0.0 ++ and self.wait_queue_entry_time > 0.0 ++ ): ++ return self.wait_queue_entry_time - self.decode_transfer_queue_entry_time ++ return None ++ ++ def get_pd_decode_forward_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.forward_entry_time > 0.0 ++ and self.completion_time > 0.0 ++ ): ++ return self.completion_time - self.forward_entry_time ++ return None ++ ++ def get_pd_bootstrap_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_bootstrap_duration is not None: ++ return self.fwd_bootstrap_duration ++ if self.bootstrap_done_time <= 0.0: ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.prefill_bootstrap_queue_entry_time > 0.0 ++ ): ++ return self.bootstrap_done_time - self.prefill_bootstrap_queue_entry_time ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_prealloc_queue_entry_time > 0.0 ++ ): ++ return self.bootstrap_done_time - self.decode_prealloc_queue_entry_time ++ return None ++ ++ def get_pd_alloc_waiting_duration(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_alloc_waiting_duration is not None: ++ return self.fwd_alloc_waiting_duration ++ if self.bootstrap_done_time <= 0.0: ++ return None ++ if ( ++ self.disagg_mode == DisaggregationMode.PREFILL ++ and self.wait_queue_entry_time > 0.0 ++ ): ++ return self.wait_queue_entry_time - self.bootstrap_done_time ++ if ( ++ self.disagg_mode == DisaggregationMode.DECODE ++ and self.decode_transfer_queue_entry_time > 0.0 ++ ): ++ return self.decode_transfer_queue_entry_time - self.bootstrap_done_time ++ return None ++ ++ def get_pd_transfer_speed_gb_s(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_transfer_speed_gb_s is not None: ++ return self.fwd_transfer_speed_gb_s ++ if ( ++ self.disagg_mode != DisaggregationMode.NULL ++ and self.transfer_speed_gb_s > 0.0 ++ ): ++ return self.transfer_speed_gb_s ++ return None ++ ++ def get_pd_transfer_total_mb(self) -> Optional[float]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_transfer_total_mb is not None: ++ return self.fwd_transfer_total_mb ++ if self.disagg_mode != DisaggregationMode.NULL and self.transfer_total_mb > 0.0: ++ return self.transfer_total_mb ++ return None ++ ++ def get_pd_prefill_retry_count(self) -> Optional[int]: ++ if not is_slime_profiling_enabled(): ++ return None ++ if self.fwd_prefill_retry_count is not None: ++ return self.fwd_prefill_retry_count ++ if self.disagg_mode == DisaggregationMode.PREFILL: ++ return self.prefill_retry_count ++ return None ++ + def convert_to_duration(self) -> str: + if self.disagg_mode == DisaggregationMode.NULL: + queue_duration = self.forward_entry_time - self.wait_queue_entry_time +@@ -1038,6 +1214,24 @@ class SchedulerReqTimeStats(ReqTimeStatsBase): + "prefill_launch_latency": self.get_prefill_launch_latency(), + } + ) ++ if is_slime_profiling_enabled(): ++ for key, value in { ++ "pd_prefill_bootstrap_queue_duration": ( ++ self.get_pd_prefill_bootstrap_queue_duration() ++ ), ++ "pd_prefill_forward_duration": self.get_pd_prefill_forward_duration(), ++ "pd_prefill_transfer_queue_duration": self.get_pd_prefill_transfer_queue_duration(), ++ "pd_decode_prealloc_duration": self.get_pd_decode_prealloc_duration(), ++ "pd_decode_transfer_duration": self.get_pd_decode_transfer_duration(), ++ "pd_decode_forward_duration": self.get_pd_decode_forward_duration(), ++ "pd_bootstrap_duration": self.get_pd_bootstrap_duration(), ++ "pd_alloc_waiting_duration": self.get_pd_alloc_waiting_duration(), ++ "pd_transfer_speed_gb_s": self.get_pd_transfer_speed_gb_s(), ++ "pd_transfer_total_mb": self.get_pd_transfer_total_mb(), ++ "pd_prefill_retry_count": self.get_pd_prefill_retry_count(), ++ }.items(): ++ if value is not None: ++ meta_data[key] = value + return meta_data + + def format_duration(self, duration: float) -> str: +diff --git a/python/sglang/srt/observability/scheduler_metrics_mixin.py b/python/sglang/srt/observability/scheduler_metrics_mixin.py +index ff5695ce2e..588379a85d 100644 +--- a/python/sglang/srt/observability/scheduler_metrics_mixin.py ++++ b/python/sglang/srt/observability/scheduler_metrics_mixin.py +@@ -883,12 +883,42 @@ class SchedulerMetricsMixin: + num_tokens += sum(req.seqlen for queue in waiting_queues for req in queue) + num_waiting_reqs = sum(len(queue) for queue in waiting_queues) + ++ queue_names = ["waiting_queue"] ++ if self.disaggregation_mode == DisaggregationMode.PREFILL: ++ queue_names.append("bootstrap_queue") ++ elif self.disaggregation_mode == DisaggregationMode.DECODE: ++ queue_names.append("prealloc_queue") ++ queue_names.append("transfer_queue") ++ queue_names.append("retracted_queue") ++ ++ queue_details = [] ++ for name, queue in zip(queue_names, waiting_queues): ++ reqs_info = [{"seqlen": req.seqlen} for req in queue] ++ queue_details.append( ++ { ++ "name": name, ++ "num_reqs": len(queue), ++ "num_tokens": sum(req_info["seqlen"] for req_info in reqs_info), ++ "reqs": reqs_info, ++ } ++ ) ++ ++ running_reqs_info = [ ++ {"seqlen": req.seqlen} for req in self.running_batch.reqs ++ ] ++ running_details = { ++ "num_reqs": len(self.running_batch.reqs), ++ "reqs": running_reqs_info, ++ } ++ + return GetLoadReqOutput( + dp_rank=self.dp_rank, + num_reqs=len(self.running_batch.reqs) + num_waiting_reqs, + num_waiting_reqs=num_waiting_reqs, + num_tokens=num_tokens, + ts_tic=time.perf_counter(), ++ queue_details=queue_details, ++ running_details=running_details, + ) + + def get_loads(self: Scheduler, req: GetLoadsReqInput = None) -> GetLoadsReqOutput: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py -index b080aeb168..5b29ebf566 100644 +index d91ced805f..4c8774bb64 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py -@@ -635,6 +635,7 @@ class ServerArgs: +@@ -670,6 +670,7 @@ class ServerArgs: # Context parallelism used in the long sequence prefill phase of DeepSeek v3.2 enable_nsa_prefill_context_parallel: bool = False nsa_prefill_cp_mode: str = "round-robin-split" + disable_indexer_rope_neox_style: bool = False enable_fused_qk_norm_rope: bool = False enable_precise_embedding_interpolation: bool = False - -@@ -4781,6 +4782,12 @@ class ServerArgs: + enable_fused_moe_sum_all_reduce: bool = False +@@ -5659,6 +5660,12 @@ class ServerArgs: help="Token splitting mode for the prefill phase of DeepSeek v3.2 under context parallelism. Optional values: 'round-robin-split'(default), 'in-seq-split' " "'round-robin-split' distributes tokens across ranks based on token_idx %% cp_size. It supports multi-batch prefill, fused MoE, and FP8 KV cache.", ) @@ -3111,44 +2226,43 @@ index b080aeb168..5b29ebf566 100644 + "If the environment variable INDEXER_ROPE_NEOX_STYLE is also set and conflicts, an error is raised.", + ) parser.add_argument( - "--enable-fused-qk-norm-rope", + "--enable-prefill-context-parallel", action="store_true", diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -index 5fe45086ca..b283d2e9bd 100644 +index 40e859b2d6..2604ae037c 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py -@@ -341,7 +341,10 @@ class EAGLEDraftCudaGraphRunner: - self.seq_lens.fill_(self.seq_len_fill_value) - self.out_cache_loc.zero_() - self.positions.zero_() -- -+ self.topk_p.zero_() -+ self.topk_index.zero_() -+ self.hidden_states.zero_() -+ self.req_pool_indices.zero_() +@@ -377,6 +377,10 @@ class EAGLEDraftCudaGraphRunner: + buffers.seq_lens.fill_(self.seq_len_fill_value) + buffers.out_cache_loc.zero_() + buffers.positions.zero_() ++ buffers.topk_p.zero_() ++ buffers.topk_index.zero_() ++ buffers.hidden_states.zero_() ++ buffers.req_pool_indices.zero_() + num_tokens = bs * self.num_tokens_per_bs - # Common inputs -@@ -350,8 +353,12 @@ class EAGLEDraftCudaGraphRunner: +@@ -386,8 +390,12 @@ class EAGLEDraftCudaGraphRunner: forward_batch.out_cache_loc ) - self.positions[:raw_num_token].copy_(forward_batch.positions) -- self.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p) -- self.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index) -+ self.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p.clamp(0, 1)) -+ self.topk_index[:raw_bs].copy_( + buffers.positions[:raw_num_token].copy_(forward_batch.positions) +- buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p) +- buffers.topk_index[:raw_bs].copy_(forward_batch.spec_info.topk_index) ++ buffers.topk_p[:raw_bs].copy_(forward_batch.spec_info.topk_p.clamp(0, 1)) ++ buffers.topk_index[:raw_bs].copy_( + forward_batch.spec_info.topk_index.clamp( + 0, self.model_runner.model_config.vocab_size - 1 + ) + ) - self.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states) - self.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) + buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states) + buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices) diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py -index ac629c7ee5..c039d23508 100644 +index dbb91f555e..a04caefc34 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py -@@ -774,6 +774,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): +@@ -776,6 +776,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): self.topk_index = self.topk_index[: len(new_indices)] self.hidden_states = self.hidden_states[: len(new_indices)] self.verified_id = self.verified_id[: len(new_indices)] @@ -3159,7 +2273,7 @@ index ac629c7ee5..c039d23508 100644 else: # in some cases(e.g draft_extend), we have not filtered the batch by `unfinished_index` self.topk_p = self.topk_p[new_indices] -@@ -805,6 +809,27 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): +@@ -807,6 +811,27 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin): self.verified_id = torch.cat([self.verified_id, spec_info.verified_id], axis=0) self.topk_p = torch.cat([self.topk_p, spec_info.topk_p]) self.topk_index = torch.cat([self.topk_index, spec_info.topk_index]) @@ -3188,14 +2302,13 @@ index ac629c7ee5..c039d23508 100644 @dataclass diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py -index 4636128fa7..a9b61df393 100644 +index b0be70d751..44a78d684e 100644 --- a/python/sglang/srt/utils/common.py +++ b/python/sglang/srt/utils/common.py -@@ -2359,6 +2359,8 @@ class SafeUnpickler(pickle.Unpickler): - "sglang.srt.model_executor.model_runner.", +@@ -2157,6 +2157,7 @@ class SafeUnpickler(pickle.Unpickler): "sglang.srt.layers.", "sglang.srt.utils.", -+ # --- slime --- + "torch_npu.", + "slime.", } diff --git a/docker/version.txt b/docker/version.txt index 69e1776f8d..ba2be939c8 100644 --- a/docker/version.txt +++ b/docker/version.txt @@ -1 +1 @@ -nightly-dev-20260428a +nightly-dev-20260430a diff --git a/docs/en/developer_guide/ci.md b/docs/en/developer_guide/ci.md index 7e3b80968a..1ffaac40ec 100644 --- a/docs/en/developer_guide/ci.md +++ b/docs/en/developer_guide/ci.md @@ -75,7 +75,7 @@ NUM_GPUS = 4 # This constant is used by run-ci-changed def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") # Download datasets as needed ... def execute(): diff --git a/docs/zh/developer_guide/ci.md b/docs/zh/developer_guide/ci.md index a8e78e1b50..603b9a1734 100644 --- a/docs/zh/developer_guide/ci.md +++ b/docs/zh/developer_guide/ci.md @@ -75,7 +75,7 @@ NUM_GPUS = 4 # 此常量会被 run-ci-changed 自动读取 def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") # 按需下载数据集 ... def execute(): diff --git a/examples/eval_multi_task/multi_task.yaml b/examples/eval_multi_task/multi_task.yaml index 83ae67f8ac..bad2d61412 100644 --- a/examples/eval_multi_task/multi_task.yaml +++ b/examples/eval_multi_task/multi_task.yaml @@ -7,11 +7,11 @@ eval: path: /root/aime-2024/aime-2024.jsonl rm_type: deepscaler n_samples_per_eval_prompt: 16 - - name: gpqa # huggingface-cli download --repo-type dataset zyzshishui0627/gpqa_diamond --local-dir /root/gpqa + - name: gpqa # hf download --repo-type dataset zyzshishui0627/gpqa_diamond --local-dir /root/gpqa path: /root/gpqa/gpqa_eval.jsonl rm_type: gpqa n_samples_per_eval_prompt: 2 - - name: ifbench # huggingface-cli download --repo-type dataset zyzshishui0627/IFBench --local-dir /root/ifbench + - name: ifbench # hf download --repo-type dataset zyzshishui0627/IFBench --local-dir /root/ifbench path: /root/ifbench/IFBench_eval.jsonl rm_type: ifbench n_samples_per_eval_prompt: 1 diff --git a/examples/search-r1/README_zh.md b/examples/search-r1/README_zh.md index a0d273a278..f7d9fbb0d9 100644 --- a/examples/search-r1/README_zh.md +++ b/examples/search-r1/README_zh.md @@ -47,7 +47,7 @@ python $WORK_DIR/scripts/data_process/qa_search_test_merge.py \ ```bash # hf checkpoint -huggingface-cli download Qwen/Qwen2.5-3B --local-dir /root/Qwen2.5-3B +hf download Qwen/Qwen2.5-3B --local-dir /root/Qwen2.5-3B # mcore checkpoint cd /root/slime diff --git a/examples/strands_sglang/README.md b/examples/strands_sglang/README.md index 0101fc0e69..1d5864a47d 100644 --- a/examples/strands_sglang/README.md +++ b/examples/strands_sglang/README.md @@ -32,7 +32,7 @@ This example connects `slime` with [`strands-sglang`](https://github.com/horizon ```bash # hf checkpoint -huggingface-cli download Qwen/Qwen3-8B --local-dir /root/models/Qwen/Qwen3-8B +hf download Qwen/Qwen3-8B --local-dir /root/models/Qwen/Qwen3-8B # mcore checkpoint cd /root/slime diff --git a/examples/tau-bench/README.md b/examples/tau-bench/README.md index 3e52d79e93..07724f3bfb 100644 --- a/examples/tau-bench/README.md +++ b/examples/tau-bench/README.md @@ -29,7 +29,7 @@ Initialize the Qwen3-4B-Instruct-2507 model needed for tool use: ```bash # hf checkpoint -huggingface-cli download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507 +hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507 # mcore checkpoint cd /root/slime diff --git a/scripts/models/qwen3.5-0.8B.sh b/scripts/models/qwen3.5-0.8B.sh new file mode 100644 index 0000000000..c235c4c1f6 --- /dev/null +++ b/scripts/models/qwen3.5-0.8B.sh @@ -0,0 +1,28 @@ +MODEL_ARGS=( + --spec "slime_plugins.models.qwen3_5" "get_qwen3_5_spec" + + --disable-bias-linear + --qk-layernorm + --group-query-attention + --num-attention-heads 8 + --num-query-groups 2 + --kv-channels 256 + --num-layers 24 + --hidden-size 1024 + --ffn-hidden-size 3584 + --use-gated-attention + + --normalization RMSNorm + --apply-layernorm-1p + --position-embedding-type rope + --norm-epsilon 1e-6 + --rotary-percent 0.25 + --swiglu + --vocab-size 248320 + + --rotary-base 10000000 + + # qwen3.5 specific + --attention-output-gate +) + diff --git a/slime/backends/megatron_utils/checkpoint.py b/slime/backends/megatron_utils/checkpoint.py index 5fc7896152..d4eb236936 100644 --- a/slime/backends/megatron_utils/checkpoint.py +++ b/slime/backends/megatron_utils/checkpoint.py @@ -135,7 +135,9 @@ def _load_checkpoint_hf(ddp_model, optimizer, args, load_path: str): logger.info(f"Load checkpoint from HuggingFace model into Megatron (path={load_path})") with megatron_bridge_utils.patch_megatron_model(ddp_model): - bridge = AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True) + bridge = megatron_bridge_utils.patch_auto_bridge_hf_config( + AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True) + ) bridge.load_hf_weights(ddp_model) # Copied from Megatron-core :: load_checkpoint (with simplifications) diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index 852720805b..bbbe962a5a 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -728,14 +728,14 @@ def save_hf_model(args, rollout_id: int, model: Sequence[DDP]) -> None: try: from megatron.bridge import AutoBridge - from slime.utils.megatron_bridge_utils import patch_megatron_model + from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_megatron_model path = Path(args.save_hf.format(rollout_id=rollout_id)) if should_log: logger.info(f"Saving model in HuggingFace format to {path}") - bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) + bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) path.mkdir(parents=True, exist_ok=True) diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py index d98fc7c005..68e62bb8d7 100644 --- a/slime/backends/megatron_utils/model_provider.py +++ b/slime/backends/megatron_utils/model_provider.py @@ -17,6 +17,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.training.arguments import core_transformer_config_from_args +from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config from slime.utils.misc import load_function @@ -85,7 +86,7 @@ def wrapped_model_provider( import slime_plugins.megatron_bridge # noqa: F401 # register custom bridges - bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True) + bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)) provider = bridge.to_megatron_provider(load_weights=False) # TODO: we should not manually set this... provider.tensor_model_parallel_size = args.tensor_model_parallel_size diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py index 638d8fd1a9..9ea728697b 100644 --- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py +++ b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py @@ -44,7 +44,9 @@ def __init__(self, *args, **kwargs): import slime_plugins.megatron_bridge # noqa: F401 - self._bridge = AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True) + self._bridge = megatron_bridge_utils.patch_auto_bridge_hf_config( + AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True) + ) _patch_bridge_expert_cache_to_cpu() def get_hf_weight_chunks(self, megatron_local_weights): diff --git a/slime/utils/megatron_bridge_utils.py b/slime/utils/megatron_bridge_utils.py index 9e5f065cd4..c87fb5b7b0 100644 --- a/slime/utils/megatron_bridge_utils.py +++ b/slime/utils/megatron_bridge_utils.py @@ -6,6 +6,38 @@ unwrap_model = None +def patch_hf_config_for_megatron_bridge(hf_config): + configs = [] + seen_config_ids = set() + + def add_config(config): + if config is None or id(config) in seen_config_ids: + return + seen_config_ids.add(id(config)) + configs.append(config) + + add_config(hf_config) + add_config(getattr(hf_config, "config", None)) + + for config in list(configs): + add_config(getattr(config, "text_config", None)) + + for config in configs: + rope_params = getattr(config, "rope_parameters", None) or getattr(config, "rope_scaling", None) + if isinstance(rope_params, dict) and "rope_theta" in rope_params and not hasattr(config, "rope_theta"): + config.rope_theta = rope_params["rope_theta"] + + return hf_config + + +def patch_auto_bridge_hf_config(bridge): + hf_pretrained = getattr(bridge, "hf_pretrained", None) + if hf_pretrained is not None: + patch_hf_config_for_megatron_bridge(hf_pretrained) + + return bridge + + @contextmanager def patch_megatron_model(model): unwrapped_model = unwrap_model(model)[0] diff --git a/slime_plugins/mbridge/qwen3_5.py b/slime_plugins/mbridge/qwen3_5.py index 4cdc54691b..d01094ecc9 100644 --- a/slime_plugins/mbridge/qwen3_5.py +++ b/slime_plugins/mbridge/qwen3_5.py @@ -117,6 +117,15 @@ def _get_text_config(self): return self.hf_config.text_config return self.hf_config + def _adjust_mapping_for_shared_weights(self): + text_config = self._get_text_config() + tie_word_embeddings = getattr(text_config, "tie_word_embeddings", False) or getattr( + self.hf_config, "tie_word_embeddings", False + ) + if tie_word_embeddings: + self._DIRECT_MAPPING = dict(self._DIRECT_MAPPING) + self._DIRECT_MAPPING["output_layer.weight"] = "model.language_model.embed_tokens.weight" + def _supports_transformer_config_kwarg(self, kwarg_name: str) -> bool: """Check whether the current TransformerConfig accepts a given kwarg.""" transformer_config_class = getattr(self, "TransformerConfigClass", None) diff --git a/tests/test_glm4.7_30B_A3B_pd_mooncake.py b/tests/test_glm4.7_30B_A3B_pd_mooncake.py new file mode 100644 index 0000000000..7048c34c69 --- /dev/null +++ b/tests/test_glm4.7_30B_A3B_pd_mooncake.py @@ -0,0 +1,159 @@ +"""GLM-4.7-Flash colocated training test with single-node PD + Mooncake.""" + +import os +import tempfile + +import yaml + +import slime.utils.external_utils.command_utils as U + + +MODEL_REPO = "zai-org/GLM-4.7-Flash" +MODEL_NAME = "GLM-4.7-Flash" +MODEL_TYPE = "glm4.7-30B-A3B" +NUM_GPUS = 8 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download {MODEL_REPO} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/root/models", + hf_checkpoint=f"/root/models/{MODEL_NAME}", + ) + + +def write_sglang_config() -> str: + config = { + "sglang": [ + { + "name": "default", + "server_groups": [ + { + "worker_type": "prefill", + "num_gpus": 4, + "num_gpus_per_engine": 4, + "overrides": {"disaggregation_transfer_backend": "mooncake"}, + }, + { + "worker_type": "decode", + "num_gpus": 4, + "num_gpus_per_engine": 4, + "overrides": {"disaggregation_transfer_backend": "mooncake"}, + }, + ], + } + ] + } + f = tempfile.NamedTemporaryFile("w", suffix=".yaml", prefix="sglang_pd_mooncake_", delete=False) + with f: + yaml.safe_dump(config, f, sort_keys=False) + return f.name + + +def execute(): + sglang_config = write_sglang_config() + + ckpt_args = ( + f"--hf-checkpoint /root/models/{MODEL_NAME} " + f"--ref-load /root/models/{MODEL_NAME}_torch_dist " + ) + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 2 " + "--rollout-max-response-len 512 " + "--rollout-temperature 0.8 " + "--global-batch-size 8 " + ) + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + "--eps-clip-high 0.28 " + ) + perf_args = ( + "--tensor-model-parallel-size 2 " + "--sequence-parallel " + "--pipeline-model-parallel-size 2 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 4 " + "--expert-tensor-parallel-size 1 " + "--decoder-last-pipeline-num-layers 23 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 2048 " + ) + sglang_args = ( + "--rollout-num-gpus 8 " + "--rollout-num-gpus-per-engine 4 " + "--sglang-enable-dp-attention " + "--sglang-dp-size 4 " + "--sglang-enable-dp-lm-head " + "--sglang-ep-size 4 " + "--sglang-moe-dense-tp-size 1 " + "--sglang-mem-fraction-static 0.45 " + "--sglang-cuda-graph-max-bs 8 " + "--sglang-max-running-requests 16 " + "--sglang-disaggregation-transfer-backend mooncake " + "--sglang-watchdog-timeout 1200 " + "--sglang-router-request-timeout-secs 1200 " + "--sglang-enable-metrics " + f"--sglang-config {sglang_config} " + ) + misc_args = ( + "--ci-test " + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--moe-token-dispatcher-type alltoall " + ) + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{sglang_args} " + f"{misc_args} " + ) + U.execute_train(train_args=train_args, num_gpus_per_node=NUM_GPUS, megatron_model_type=MODEL_TYPE) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen2.5_0.5B_async_short.py b/tests/test_qwen2.5_0.5B_async_short.py index 43a8f0a156..c3925d4432 100644 --- a/tests/test_qwen2.5_0.5B_async_short.py +++ b/tests/test_qwen2.5_0.5B_async_short.py @@ -10,7 +10,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/dapo-math-17k") diff --git a/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py b/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py index 63fd24904a..6e687799c8 100644 --- a/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py +++ b/tests/test_qwen2.5_0.5B_debug_rollout_then_train.py @@ -23,7 +23,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/test_qwen2.5_0.5B_opd_sglang.py b/tests/test_qwen2.5_0.5B_opd_sglang.py index 85addd0ab4..eb0892fb91 100644 --- a/tests/test_qwen2.5_0.5B_opd_sglang.py +++ b/tests/test_qwen2.5_0.5B_opd_sglang.py @@ -18,7 +18,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py b/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py index 515592ab7c..c03d863d8e 100644 --- a/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py +++ b/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py @@ -12,7 +12,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/dapo-math-17k") diff --git a/tests/test_qwen2.5_0.5B_sglang_config.py b/tests/test_qwen2.5_0.5B_sglang_config.py index 146ca2f18f..f30c6ac6cf 100644 --- a/tests/test_qwen2.5_0.5B_sglang_config.py +++ b/tests/test_qwen2.5_0.5B_sglang_config.py @@ -30,7 +30,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/test_qwen2.5_0.5B_sglang_config_distributed.py b/tests/test_qwen2.5_0.5B_sglang_config_distributed.py index a8306f10e2..68215b34c6 100644 --- a/tests/test_qwen2.5_0.5B_sglang_config_distributed.py +++ b/tests/test_qwen2.5_0.5B_sglang_config_distributed.py @@ -31,7 +31,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/test_qwen2.5_0.5B_short.py b/tests/test_qwen2.5_0.5B_short.py index c5b3157848..6f45095bfb 100644 --- a/tests/test_qwen2.5_0.5B_short.py +++ b/tests/test_qwen2.5_0.5B_short.py @@ -10,7 +10,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/dapo-math-17k") diff --git a/tests/test_qwen2.5_0.5B_gsm8k_async_short.py b/tests/test_qwen3.5_0.8B_gsm8k_async_short.py similarity index 85% rename from tests/test_qwen2.5_0.5B_gsm8k_async_short.py rename to tests/test_qwen3.5_0.8B_gsm8k_async_short.py index ee71ff60e6..14d92ff4fe 100644 --- a/tests/test_qwen2.5_0.5B_gsm8k_async_short.py +++ b/tests/test_qwen3.5_0.8B_gsm8k_async_short.py @@ -1,21 +1,30 @@ import os + import slime.utils.external_utils.command_utils as U + TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" +MODEL_NAME = "Qwen3.5-0.8B" +MODEL_TYPE = "qwen3.5-0.8B" NUM_GPUS = 4 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load {TORCH_DIST_CKPT} " rollout_args = ( "--prompt-data /root/datasets/gsm8k/train.parquet " @@ -94,10 +103,10 @@ def execute(): "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " + "--loss-mask-type qwen3_5 " "--actor-num-nodes 1 " "--actor-num-gpus-per-node 1 " "--rollout-num-gpus 3 " - "--megatron-to-hf-mode bridge " ) train_args = ( @@ -124,8 +133,6 @@ def execute(): if __name__ == "__main__": prepare() - os.environ.pop("http_proxy") - os.environ.pop("https_proxy") - os.environ.pop("HTTP_PROXY") - os.environ.pop("HTTPS_PROXY") + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) execute() diff --git a/tests/test_qwen2.5_0.5B_gsm8k_short.py b/tests/test_qwen3.5_0.8B_gsm8k_short.py similarity index 85% rename from tests/test_qwen2.5_0.5B_gsm8k_short.py rename to tests/test_qwen3.5_0.8B_gsm8k_short.py index 1b47c8d007..856413de75 100644 --- a/tests/test_qwen2.5_0.5B_gsm8k_short.py +++ b/tests/test_qwen3.5_0.8B_gsm8k_short.py @@ -1,21 +1,30 @@ import os + import slime.utils.external_utils.command_utils as U + TIGHT_DEVICE_MEMORY = U.get_bool_env_var("SLIME_TEST_TIGHT_DEVICE_MEMORY", "1") -MODEL_NAME = "Qwen2.5-0.5B-Instruct" -MODEL_TYPE = "qwen2.5-0.5B" +MODEL_NAME = "Qwen3.5-0.8B" +MODEL_TYPE = "qwen3.5-0.8B" NUM_GPUS = 4 +TORCH_DIST_CKPT = f"/dev/shm/{MODEL_NAME}_torch_dist" def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + dir_dst="/dev/shm", + ) def execute(): - ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ " + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load {TORCH_DIST_CKPT} " rollout_args = ( "--prompt-data /root/datasets/gsm8k/train.parquet " @@ -94,10 +103,10 @@ def execute(): "--accumulate-allreduce-grads-in-fp32 " "--attention-softmax-in-fp32 " "--attention-backend flash " + "--loss-mask-type qwen3_5 " "--actor-num-nodes 1 " "--actor-num-gpus-per-node 4 " "--colocate " - "--megatron-to-hf-mode bridge " ) train_args = ( @@ -123,8 +132,6 @@ def execute(): if __name__ == "__main__": prepare() - os.environ.pop("http_proxy") - os.environ.pop("https_proxy") - os.environ.pop("HTTP_PROXY") - os.environ.pop("HTTPS_PROXY") + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) execute() diff --git a/tests/test_qwen3_30B_A3B_pd_mooncake.py b/tests/test_qwen3_30B_A3B_pd_mooncake.py new file mode 100644 index 0000000000..7ff1ed652f --- /dev/null +++ b/tests/test_qwen3_30B_A3B_pd_mooncake.py @@ -0,0 +1,145 @@ +import os +import tempfile + +import slime.utils.external_utils.command_utils as U + + +MODEL_NAME = "Qwen3-30B-A3B" +MODEL_TYPE = "qwen3-30B-A3B" +NUM_GPUS = 8 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + U.hf_download_dataset("zhuzilin/aime-2024") + U.convert_checkpoint( + model_name=MODEL_NAME, + megatron_model_type=MODEL_TYPE, + num_gpus_per_node=NUM_GPUS, + ) + + +def execute(): + debug_data_path = os.environ.get("DEBUG_ROLLOUT_DATA") or tempfile.mktemp( + prefix="qwen3_30b_a3b_pd_rollout_", suffix=".pt" + ) + try: + os.remove(debug_data_path) + except FileNotFoundError: + pass + print(f"Saving debug rollout data to {debug_data_path}") + + ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} " f"--ref-load /root/{MODEL_NAME}_torch_dist " + + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 8 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 16384 " + "--rollout-temperature 1.0 " + "--global-batch-size 32 " + ) + + eval_args = ( + "--eval-prompt-data aime24 /root/datasets/aime-2024/aime-2024.jsonl " + "--n-samples-per-eval-prompt 2 " + "--eval-max-response-len 16384 " + "--eval-temperature 0.6 " + "--eval-top-p 0.95 " + ) + + perf_args = ( + "--tensor-model-parallel-size 4 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 2 " + "--expert-model-parallel-size 8 " + "--expert-tensor-parallel-size 1 " + "--recompute-granularity full " + "--recompute-method uniform " + "--recompute-num-layers 1 " + "--max-tokens-per-gpu 8192 " + ) + + grpo_args = ( + "--advantage-estimator grpo " + "--kl-loss-coef 0.00 " + "--kl-loss-type low_var_kl " + "--kl-coef 0.00 " + "--entropy-coef 0.00 " + "--eps-clip 0.2 " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + "--optimizer-cpu-offload " + "--overlap-cpu-optimizer-d2h-h2d " + "--use-precision-aware-optimizer " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 4 " + "--sglang-mem-fraction-static 0.75 " + "--sglang-enable-dp-attention " + "--sglang-dp-size 4 " + "--sglang-ep-size 4 " + "--sglang-enable-dp-lm-head " + "--sglang-cuda-graph-bs 1 2 4 8 16 24 32 " + "--sglang-max-running-requests 512 " + "--prefill-num-servers 1 " + "--sglang-enable-metrics " + ) + + misc_args = ( + "--ci-test " + f"--save-debug-rollout-data {debug_data_path} " + "--update-weight-buffer-size 2147483648 " + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + "--actor-num-gpus-per-node 8 " + "--colocate " + "--moe-token-dispatcher-type flex " + "--moe-enable-deepep " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{grpo_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{eval_args} " + f"{sglang_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute() diff --git a/tests/test_qwen3_5_mtp_bridge_mapping.py b/tests/test_qwen3_5_mtp_bridge_mapping.py index c8a970ef29..cb304ab5db 100644 --- a/tests/test_qwen3_5_mtp_bridge_mapping.py +++ b/tests/test_qwen3_5_mtp_bridge_mapping.py @@ -168,6 +168,18 @@ def test_mtp_block_spec_uses_current_transformer_layer_spec(): assert result["mtp_block_spec"] == ("mtp-spec", "REAL_LAYER_SPEC_VP3") +@pytest.mark.unit +def test_tied_qwen3_5_uses_language_embedding_for_output_layer(): + module = load_bridge_module() + bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge) + bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(tie_word_embeddings=True)) + + bridge._adjust_mapping_for_shared_weights() + + assert bridge._DIRECT_MAPPING["output_layer.weight"] == "model.language_model.embed_tokens.weight" + assert module.Qwen3_5Bridge._DIRECT_MAPPING["output_layer.weight"] == "lm_head.weight" + + @pytest.mark.unit def test_eh_proj_keeps_column_order_when_loading_to_mcore(): module = load_bridge_module() diff --git a/tests/test_sglang_config_mixed_offload.py b/tests/test_sglang_config_mixed_offload.py index 1581a0302c..90d3d97389 100644 --- a/tests/test_sglang_config_mixed_offload.py +++ b/tests/test_sglang_config_mixed_offload.py @@ -44,7 +44,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/test_sglang_config_mixed_offload_ft.py b/tests/test_sglang_config_mixed_offload_ft.py index 8387bb62c7..f017965c5f 100644 --- a/tests/test_sglang_config_mixed_offload_ft.py +++ b/tests/test_sglang_config_mixed_offload_ft.py @@ -41,7 +41,7 @@ def prepare(): U.exec_command("mkdir -p /root/models /root/datasets") - U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") U.hf_download_dataset("zhuzilin/gsm8k") diff --git a/tests/utils/test_megatron_bridge_utils.py b/tests/utils/test_megatron_bridge_utils.py new file mode 100644 index 0000000000..d47babcaca --- /dev/null +++ b/tests/utils/test_megatron_bridge_utils.py @@ -0,0 +1,64 @@ +import types + +import pytest + +from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config, patch_hf_config_for_megatron_bridge + + +@pytest.mark.unit +def test_patch_hf_config_adds_rope_theta_from_rope_parameters(): + hf_config = types.SimpleNamespace(rope_parameters={"rope_theta": 1000000}) + + patched_config = patch_hf_config_for_megatron_bridge(hf_config) + + assert patched_config is hf_config + assert hf_config.rope_theta == 1000000 + + +@pytest.mark.unit +def test_patch_hf_config_does_not_override_existing_rope_theta(): + hf_config = types.SimpleNamespace(rope_theta=500000, rope_parameters={"rope_theta": 1000000}) + + patch_hf_config_for_megatron_bridge(hf_config) + + assert hf_config.rope_theta == 500000 + + +@pytest.mark.unit +def test_patch_hf_config_handles_nested_text_config(): + text_config = types.SimpleNamespace(rope_parameters={"rope_theta": 10000}) + hf_config = types.SimpleNamespace(text_config=text_config) + + patch_hf_config_for_megatron_bridge(hf_config) + + assert text_config.rope_theta == 10000 + + +@pytest.mark.unit +def test_patch_hf_config_handles_pretrained_wrapper_config(): + wrapped_config = types.SimpleNamespace(rope_parameters={"rope_theta": 10000}) + hf_pretrained = types.SimpleNamespace(config=wrapped_config) + + patch_hf_config_for_megatron_bridge(hf_pretrained) + + assert wrapped_config.rope_theta == 10000 + + +@pytest.mark.unit +def test_patch_hf_config_uses_rope_scaling_fallback(): + hf_config = types.SimpleNamespace(rope_scaling={"rope_theta": 10000}) + + patch_hf_config_for_megatron_bridge(hf_config) + + assert hf_config.rope_theta == 10000 + + +@pytest.mark.unit +def test_patch_auto_bridge_hf_config_patches_hf_pretrained(): + hf_config = types.SimpleNamespace(rope_parameters={"rope_theta": 12345}) + bridge = types.SimpleNamespace(hf_pretrained=hf_config) + + patched_bridge = patch_auto_bridge_hf_config(bridge) + + assert patched_bridge is bridge + assert bridge.hf_pretrained.rope_theta == 12345 diff --git a/tools/convert_torch_dist_to_hf_bridge.py b/tools/convert_torch_dist_to_hf_bridge.py index 3412e73e69..798503f218 100644 --- a/tools/convert_torch_dist_to_hf_bridge.py +++ b/tools/convert_torch_dist_to_hf_bridge.py @@ -4,6 +4,8 @@ import megatron.bridge.training.model_load_save as _model_load_save_module from megatron.bridge import AutoBridge +from slime.utils.megatron_bridge_utils import patch_auto_bridge_hf_config + # Here we need to patch Megatron Bridge's `load_model_config`, since the checkpoint is saved # by Megatron and lack of provider information. @@ -49,7 +51,7 @@ def _patched_load_model_config(checkpoint_path): raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.") print(f"Loading config from {args.origin_hf_dir}") - bridge = AutoBridge.from_hf_pretrained(args.origin_hf_dir, trust_remote_code=True) + bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.origin_hf_dir, trust_remote_code=True)) # Use Bridge's provider so the correct model class is created (e.g., Qwen3VLModel # instead of GPTModel). This is needed because MLM checkpoints lack run_config.yaml.