diff --git a/.github/workflows/conda-ci.yml b/.github/workflows/conda-ci.yml
index 6cb0f5ce2b..8fdedb8f7f 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: --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
+      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
 
     defaults:
       run:
diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml
index 8cb815c1f2..7aebedc97b 100644
--- a/.github/workflows/pr-test.yml
+++ b/.github/workflows/pr-test.yml
@@ -29,13 +29,28 @@ jobs:
 
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-short'))
 
-
     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_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"}]
+        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_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -51,74 +66,49 @@ 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
-            '
-
-
-  e2e-test-sglang-config:
-
-    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-sglang-config'))
+  e2e-test-fsdp:
 
+    if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-fsdp'))
 
     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_qwen2.5_0.5B_sglang_config.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_sglang_config_distributed.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload.py"}, {"num_gpus": 8, "test_file": "test_sglang_config_mixed_offload_ft.py"}]
+        info: [{"num_gpus": 4, "test_file": "test_qwen3_4B_fsdp_true_on_policy.py --colocated"}, {"num_gpus": 8, "test_file": "test_qwen3_vl_4B_fsdp.py"}, {"num_gpus": 4, "test_file": "test_qwen3_0.6B_megatron_fsdp_align.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -134,74 +124,49 @@ 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: |
-
-          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
-            '
-
+          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
 
   e2e-test-megatron:
 
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-megatron'))
 
-
     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_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.6_35B_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_disaggregate.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_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"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -217,74 +182,49 @@ 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: |
-
-          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
-            '
-
+          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
 
   e2e-test-precision:
 
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-precision'))
 
-
     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_0.6B_parallel_check.py"}]
+        info: [{"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 4, "test_file": "test_qwen3_0.6B_megatron_fsdp_align.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -300,74 +240,49 @@ 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: |
-
-          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
-            '
-
+          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
 
   e2e-test-ckpt:
 
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, 'run-ci-ckpt'))
 
-
     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_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}]
+        info: [{"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_ckpt.py --async-save"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -383,74 +298,49 @@ 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: |
-
-          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
-            '
-
+          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
 
   e2e-test-plugin-contracts:
 
     if: github.event_name == 'pull_request' || github.event_name == 'workflow_dispatch'
 
-
-    runs-on: ubuntu-latest
-
+    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": 0, "test_file": "test_megatron_argument_validation.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
+        info: [{"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_rollout_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_runtime_hook_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_path_loading_contracts.py"}, {"num_gpus": 0, "test_file": "plugin_contracts/test_plugin_generate_contracts.py"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -466,56 +356,49 @@ jobs:
       - name: Checkout repository
         uses: actions/checkout@v4
 
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-          cache: 'pip'
-
-      - name: Install dependencies
-        shell: bash
-        run: |
-          pip install torch --index-url https://download.pytorch.org/whl/cpu
-          pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow
-
       - name: Install
         shell: bash
-        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps
-
+        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
-          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" "${TEST_ARGS_ARRAY[@]}"
+            python "$TEST_PATH"
           else
-            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
+            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH"
           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_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.6_35B_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"}]
+        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": 2, "test_file": "test_qwen3_4B_fsdp_true_on_policy.py"}, {"num_gpus": 8, "test_file": "test_qwen3_vl_4B_fsdp.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": 4, "test_file": "test_qwen3_0.6B_megatron_fsdp_align.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"}]
     defaults:
       run:
         working-directory: ${{ github.workspace }}
@@ -531,67 +414,37 @@ 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: |
-
-          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
-            '
-
+          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
 
 
   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 }}
@@ -632,6 +485,23 @@ 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) }}
@@ -650,55 +520,19 @@ 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: |
-          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
+          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
diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2
index 50fe472273..dd211bd8ae 100644
--- a/.github/workflows/pr-test.yml.j2
+++ b/.github/workflows/pr-test.yml.j2
@@ -2,31 +2,28 @@
     'e2e-test-short': {
       'label': 'run-ci-short',
       'tests': [
-        {'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},
+        {'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_sglang_config.py', 'num_gpus': 8},
+        {'test_file': 'test_qwen2.5_0.5B_sglang_config_distributed.py', 'num_gpus': 8},
       ],
     },
-    'e2e-test-sglang-config': {
-      'label': 'run-ci-sglang-config',
+    'e2e-test-fsdp': {
+      'label': 'run-ci-fsdp',
       'tests': [
-        {'test_file': 'test_qwen2.5_0.5B_sglang_config.py', 'num_gpus': 8},
-        {'test_file': 'test_qwen2.5_0.5B_sglang_config_distributed.py', 'num_gpus': 8},
-        {'test_file': 'test_sglang_config_mixed_offload.py', 'num_gpus': 8},
-        {'test_file': 'test_sglang_config_mixed_offload_ft.py', 'num_gpus': 8},
+        {'test_file': 'test_qwen3_4B_fsdp_true_on_policy.py --colocated', 'num_gpus': 4},
+        {'test_file': 'test_qwen3_vl_4B_fsdp.py', 'num_gpus': 8},
+        {'test_file': 'test_qwen3_0.6B_megatron_fsdp_align.py', 'num_gpus': 4},
       ],
     },
     'e2e-test-megatron': {
       '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.6_35B_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},
-        {'test_file': 'test_qwen3_4B_ppo_disaggregate.py', 'num_gpus': 8},
         {'test_file': 'test_qwen3_4B_ppo_train_critic_only.py', 'num_gpus': 8},
         {'test_file': 'test_moonlight_16B_A3B.py', 'num_gpus': 8},
         {'test_file': 'test_moonlight_16B_A3B_r3.py', 'num_gpus': 8, 'enable_eval': '0'},
@@ -39,22 +36,21 @@
       'label': 'run-ci-precision',
       'tests': [
         {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8},
+        {'test_file': 'test_qwen3_0.6B_megatron_fsdp_align.py', 'num_gpus': 4},
       ],
     },
     'e2e-test-ckpt': {
       'label': 'run-ci-ckpt',
       'tests': [
         {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8},
-        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', 'num_gpus': 8},
+        {'test_file': 'test_qwen3_4B_ckpt.py --async-save', 'num_gpus': 8},
       ],
     },
 
     'e2e-test-plugin-contracts': {
       'label': 'run-ci-plugin-contracts',
       'always': True,
-      'cpu': True,
       'tests': [
-        {'test_file': 'test_megatron_argument_validation.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_rollout_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_runtime_hook_contracts.py', 'num_gpus': 0},
         {'test_file': 'plugin_contracts/test_plugin_path_loading_contracts.py', 'num_gpus': 0},
@@ -66,18 +62,19 @@
       'label': 'run-ci-image',
       'image': 'slimerl/slime-test:latest',
       'tests': [
-        {'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_gsm8k_async_short.py', 'num_gpus': 4},
+        {'test_file': 'test_qwen2.5_0.5B_gsm8k_short.py', 'num_gpus': 4},
+        {'test_file': 'test_qwen3_4B_fsdp_true_on_policy.py', 'num_gpus': 2},
+        {'test_file': 'test_qwen3_vl_4B_fsdp.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.6_35B_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_0.6B_megatron_fsdp_align.py', 'num_gpus': 4},
         {'test_file': 'test_qwen3_4B_ckpt.py', 'num_gpus': 8},
-        {'test_file': 'test_qwen3_4B_ckpt.py', 'test_args': '--async-save', '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', 'num_gpus': 8},
       ],
@@ -112,11 +109,24 @@ jobs:
 <% else %>
     if: (github.event_name == 'workflow_dispatch') || (github.event.pull_request && contains(github.event.pull_request.labels.*.name, '<< config.label >>'))
 <% endif %>
-<% if config.get('cpu') %>
-    runs-on: ubuntu-latest
-<% else %>
     runs-on: self-hosted
-<% endif %>
+    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
     strategy:
       fail-fast: false
       matrix:
@@ -135,101 +145,38 @@ jobs:
     steps:
       - name: Checkout repository
         uses: actions/checkout@v4
-<% if config.get('cpu') %>
-
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.10'
-          cache: 'pip'
-
-      - name: Install dependencies
-        shell: bash
-        run: |
-          pip install torch --index-url https://download.pytorch.org/whl/cpu
-          pip install pytest numpy packaging pyyaml omegaconf tqdm httpx pybase64 pylatexenc sympy aiohttp pillow
 
       - name: Install
         shell: bash
-        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps
-<% else %>
-<% endif %>
+        run: cd $GITHUB_WORKSPACE && pip install -e . --no-deps --break-system-packages
 
       - 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" "${TEST_ARGS_ARRAY[@]}"
+            python "$TEST_PATH"
           else
-            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH" "${TEST_ARGS_ARRAY[@]}"
+            python tests/ci/gpu_lock_exec.py --count ${{ matrix.info.num_gpus }} -- python "$TEST_PATH"
           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 }}
@@ -270,6 +217,23 @@ 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) }}
@@ -288,55 +252,19 @@ 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: |
-          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
-            '
+          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
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 7b8b7ba51b..222ea070b1 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -6,7 +6,7 @@ Thank you for your interest in contributing to slime! We deeply appreciate every
 
 ## Collaboration Scope
 
-slime is the RL training infrastructure behind [GLM-4.5 through GLM-5.1](https://z.ai) and a large number of internal experiments at Z.ai. We open-sourced slime because we believe the training scenarios used internally cover the majority of cutting-edge RL algorithm requirements, and we hope to provide the community with a correct and efficient large-scale RL training infrastructure.
+slime is the RL training infrastructure behind [GLM-4.5 through GLM-5](https://z.ai) and a large number of internal experiments at Z.ai. We open-sourced slime because we believe the training scenarios used internally cover the majority of cutting-edge RL algorithm requirements, and we hope to provide the community with a correct and efficient large-scale RL training infrastructure.
 
 Our goal for open-source collaboration is focused on **bug fixes** and **general-purpose large-scale RL optimizations**. We have had several successful collaborations with the community in this area, including:
 
diff --git a/README.md b/README.md
index 6d0232c344..3cd50d88cf 100644
--- a/README.md
+++ b/README.md
@@ -10,8 +10,8 @@
 1.  **High-Performance Training**: Supports efficient training in various modes by connecting Megatron with SGLang;
 2.  **Flexible Data Generation**: Enables arbitrary training data generation workflows through custom data generation interfaces and server-based engines.
 
-slime is the RL-framework behind [GLM-5.1](https://z.ai/blog/glm-5.1), [GLM-5](https://z.ai/blog/glm-5), [GLM-4.7](https://z.ai/blog/glm-4.7), [GLM-4.6](https://z.ai/blog/glm-4.6), [GLM-4.5](https://z.ai/blog/glm-4.5) and apart from models from Z.ai, we also supports the following models:
-- Qwen series (Qwen3.6, Qwen3.5, Qwen3Next, Qwen3MoE, Qwen3, Qwen2.5);
+slime is the RL-framework behind [GLM-5](https://z.ai/blog/glm-5), [GLM-4.7](https://z.ai/blog/glm-4.7), [GLM-4.6](https://z.ai/blog/glm-4.6), [GLM-4.5](https://z.ai/blog/glm-4.5) and apart from models from Z.ai, we also supports the following models:
+- Qwen3 series (Qwen3Next, Qwen3MoE, Qwen3), Qwen2.5 series;
 - DeepSeek V3 series (DeepSeek V3, V3.1, DeepSeek R1);
 - Llama 3.
 
@@ -51,14 +51,6 @@ We also provide examples for some use cases not covered in the quick start guide
 
 slime has powered several novel research projects and production systems. Here are some notable examples:
 
-### 🌈 Relax: Asynchronous RL Engine for Omni-Modal Agentic Training
-
-[**Relax**](https://github.com/redai-infra/Relax) (Reinforcement Engine Leveraging Agentic X-modality) is an omni-modal agentic RL framework open-sourced by the RedAI Infra team, built upon the slime infrastructure stack that combines Ray, Megatron-LM, and SGLang. Relax adopts a service-oriented architecture on Ray Serve with Megatron-LM and SGLang as training/inference backends. It uses [TransferQueue](https://github.com/Ascend/TransferQueue) to fully decouple Actor, Rollout, ActorFwd, Reference, and Advantage computation onto independent GPU clusters, and introduces **DCS (Distributed Checkpoint Service)** — an NCCL-broadcast weight-sync engine that streams updated Actor weights to Rollout/ActorFwd/Reference asynchronously and overlaps the transfer with the next training step, enabling fully-async training at configurable staleness. Relax supports end-to-end RL for text, vision, and audio (including Qwen3-Omni) and agentic multi-turn rollouts.
-
-### 🦞 OpenClaw-RL: Train a Personalized Clawbot Simply by Talking to It
-
-[**OpenClaw-RL**](https://github.com/Gen-Verse/OpenClaw-RL) is an RL server for personalized OpenClaw agents. It hosts the OpenClaw model and improves it from prior conversations across deployments, while slime's asynchronous RL infrastructure prevents training from interfering with API serving. It supports two automatic optimization methods: GRPO with binary feedback inferred from subsequent states, and on-policy distillation that extracts hindsight hints from later feedback for the current policy.
-
 ### ⚛️ P1: Mastering Physics Olympiads with Reinforcement Learning
 
 [**P1**](https://prime-rl.github.io/P1/) is a family of open-source physics reasoning models trained entirely through reinforcement learning. P1 leverages slime as the RL post training framework, and introduces a multi-stage RL training algorithm that progressively enhances reasoning ability through adaptive learnability adjustment and stabilization mechanisms. Enpowered by this training paradigm, P1 delivers breakthrough performance in open-source physics reasoning.
diff --git a/README_zh.md b/README_zh.md
index 0c00795dbf..5bc249c070 100644
--- a/README_zh.md
+++ b/README_zh.md
@@ -10,8 +10,8 @@
 1. **高性能训练**：通过连接 Megatron 与 SGLang，支持各种模式的高效训练；
 2. **灵活的数据生成**：通过自定义数据生成接口以及 server based engine，实现任意的数据训练数据生成流程。
 
-slime 是 [GLM-5.1](https://z.ai/blog/glm-5.1)、[GLM-5](https://z.ai/blog/glm-5)、[GLM-4.7](https://z.ai/blog/glm-4.7)、[GLM-4.6](https://z.ai/blog/glm-4.6)、[GLM-4.5](https://z.ai/blog/glm-4.5) 背后的 RL 训练框架，除此之外，slime 还支持:
-- Qwen 系列 (Qwen3.6、Qwen3.5、Qwen3Next、Qwen3MoE、Qwen3、Qwen2.5)；
+slime 是 [GLM-5](https://z.ai/blog/glm-5)、[GLM-4.7](https://z.ai/blog/glm-4.7)、[GLM-4.6](https://z.ai/blog/glm-4.6)、[GLM-4.5](https://z.ai/blog/glm-4.5) 背后的 RL 训练框架，除此之外，slime 还支持:
+- Qwen3 系列 (Qwen3Next, Qwen3MoE, Qwen3), Qwen2.5 系列；
 - DeepSeek V3 系列 (DeepSeek V3, V3.1, DeepSeek R1)；
 - Llama 3。
 
diff --git a/build_conda.sh b/build_conda.sh
index 1d8210b232..43564311ad 100644
--- a/build_conda.sh
+++ b/build_conda.sh
@@ -12,7 +12,7 @@ source ~/.bashrc
 micromamba create -n slime python=3.12 pip -c conda-forge -y
 micromamba activate slime
 export CUDA_HOME="$CONDA_PREFIX"
-export SGLANG_COMMIT="bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2"
+export SGLANG_COMMIT="24c91001cf99ba642be791e099d358f4dfe955f5"
 export MEGATRON_COMMIT="3714d81d418c9f1bca4594fc35f9e8289f652862"
 
 export BASE_DIR=${BASE_DIR:-"/root"}
@@ -22,7 +22,8 @@ cd $BASE_DIR
 micromamba install -n slime cuda cuda-nvtx cuda-nvtx-dev nccl -c nvidia/label/cuda-12.9.1 -y
 micromamba install -n slime -c conda-forge cudnn -y
 
-pip install cuda-python==12.9
+# prevent installing cuda 13.0 for sglang
+pip install cuda-python==13.1.0
 pip install torch==2.9.1 torchvision==0.24.1 torchaudio==2.9.1 --index-url https://download.pytorch.org/whl/cu129
 
 # install sglang
@@ -41,7 +42,7 @@ MAX_JOBS=64 pip -v install flash-attn==2.7.4.post1 --no-build-isolation
 
 pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
 pip install --no-build-isolation "transformer_engine[pytorch]==2.10.0"
-pip install flash-linear-attention==0.4.1
+pip install flash-linear-attention==0.4.0
 NVCC_APPEND_FLAGS="--threads 4" \
   pip -v install --disable-pip-version-check --no-cache-dir \
   --no-build-isolation \
@@ -50,7 +51,6 @@ NVCC_APPEND_FLAGS="--threads 4" \
 pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall
 pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
 pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation
-pip install https://github.com/zhuzilin/sgl-router/releases/download/v0.3.2-5f8d397/sglang_router-0.3.2-cp38-abi3-manylinux_2_28_x86_64.whl --force-reinstall
 
 # megatron
 cd $BASE_DIR
@@ -68,8 +68,7 @@ if [ ! -d "$BASE_DIR/slime" ]; then
   export SLIME_DIR=$BASE_DIR/slime
   pip install -e .
 else
-  export SLIME_DIR=$BASE_DIR/slime
-  cd $SLIME_DIR
+  export SLIME_DIR=$BASE_DIR/
   pip install -e .
 fi
 
@@ -79,6 +78,6 @@ pip install "numpy<2"
 
 # apply patch
 cd $BASE_DIR/sglang
-git apply $SLIME_DIR/docker/patch/v0.5.9/sglang.patch
+git apply $SLIME_DIR/docker/patch/v0.5.7/sglang.patch
 cd $BASE_DIR/Megatron-LM
-git apply $SLIME_DIR/docker/patch/v0.5.9/megatron.patch
+git apply $SLIME_DIR/docker/patch/v0.5.7/megatron.patch
\ No newline at end of file
diff --git a/convert_qwen2.5_ckpt.sh b/convert_qwen2.5_ckpt.sh
deleted file mode 100644
index fae03587dc..0000000000
--- a/convert_qwen2.5_ckpt.sh
+++ /dev/null
@@ -1,5 +0,0 @@
-source scripts/models/qwen2.5-0.5B.sh
-PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
-    ${MODEL_ARGS[@]} \
-    --hf-checkpoint /root/Qwen2.5-0.5B-Instruct \
-    --save /root/Qwen2.5-0.5B-Instruct_torch_dist/
diff --git a/docker/Dockerfile b/docker/Dockerfile
index 3876daaf07..710512d15f 100644
--- a/docker/Dockerfile
+++ b/docker/Dockerfile
@@ -1,10 +1,10 @@
-ARG SGLANG_IMAGE_TAG=v0.5.10.post1
+ARG SGLANG_IMAGE_TAG=v0.5.9
 FROM slimerl/sglang:${SGLANG_IMAGE_TAG} AS sglang
 
 # ======================================== Arguments =============================================
 
 ARG PATCH_VERSION=latest
-ARG MEGATRON_COMMIT=1dcf0dafa884ad52ffb243625717a3471643e087
+ARG MEGATRON_COMMIT=3714d81d418c9f1bca4594fc35f9e8289f652862
 
 ARG ENABLE_CUDA_13=0
 
@@ -54,8 +54,8 @@ RUN git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \
     cd Megatron-LM && git checkout ${MEGATRON_COMMIT} && \
     pip install -e .
 
-RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@d64a639 --no-cache-dir --force-reinstall
-RUN pip install git+https://github.com/radixark/Megatron-Bridge.git@bridge --no-deps --no-build-isolation
+RUN pip install git+https://github.com/fzyzcjy/torch_memory_saver.git@dc6876905830430b5054325fa4211ff302169c6b --no-cache-dir --force-reinstall
+RUN pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
 RUN pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation
 
 # This patch from masahi will be included in later Triton releases
@@ -64,8 +64,7 @@ RUN if [ "$ENABLE_CUDA_13" = "1" ]; then \
   fi
 
 COPY requirements.txt /tmp/requirements.txt
-RUN pip install --ignore-installed PyJWT && \
-    pip install -r /tmp/requirements.txt
+RUN 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 \
@@ -79,9 +78,6 @@ RUN pip install nvidia-cudnn-cu12==9.16.0.29
 # reinstall numpy 1.x for megatron
 RUN pip install "numpy<2"
 
-RUN pip install https://github.com/zhuzilin/sgl-router/releases/download/v0.3.2-5f8d397/sglang_router-0.3.2-cp38-abi3-manylinux_2_28_x86_64.whl --force-reinstall
-RUN python -c "import sglang_router; assert 'slime' in sglang_router.__version__"
-
 RUN rm -rf /root/.cache/pip /root/flash-attention
 
 # ====================================== Patches ============================================
diff --git a/docker/Dockerfile.gb10 b/docker/Dockerfile.gb10
deleted file mode 100644
index 471e2ce6f1..0000000000
--- a/docker/Dockerfile.gb10
+++ /dev/null
@@ -1,208 +0,0 @@
-# Dockerfile for NVIDIA DGX Spark (GB10 / sm_121a, aarch64).
-#
-# slime's stock Dockerfile targets x86_64 + H100/B200/GB200 and does not work
-# on GB10 (ptxas in CUDA 12.9 lacks sm_121a; published sgl_kernel wheels have
-# no sm_121 variant and ABI-mismatch against NGC libtorch).
-#
-# This Dockerfile rebases on the NGC vLLM container which ships CUDA 13.2,
-# PyTorch 2.11.0a0 (compiled with compute_120 for Blackwell PTX forward-compat),
-# Triton 3.6.0, and flash-attn 2.7.4.post1 — all verified working on GB10.
-#
-# Digest-pinned for reproducibility. Update the digest together with the tag.
-# Companion notes: NOTES_GB10.md (15 resolved blockers documented).
-
-ARG BASE_IMAGE=nvcr.io/nvidia/vllm:26.03-py3
-ARG BASE_DIGEST=sha256:13e327dad79e6e417f6687fec2ba76b0386d597082ec0ee003c1e964ec6ad0e7
-FROM ${BASE_IMAGE}@${BASE_DIGEST}
-
-# =============================================================================
-# Common env: GB10 reports cc 12.1 but NGC torch only lists compute_120. PTX
-# forward-compat carries sm_120 PTX → sm_121 at runtime. Set the same for all
-# downstream compilation layers.
-# =============================================================================
-ENV TORCH_CUDA_ARCH_LIST="12.0+PTX"
-
-# =============================================================================
-# System packages missing from the NGC base but needed downstream:
-#  - libz3-dev : tilelang's Z3 SMT autoscheduler dlopens libz3.so
-# =============================================================================
-RUN apt-get update && apt-get install -y --no-install-recommends \
-        libz3-dev \
-    && rm -rf /var/lib/apt/lists/*
-
-# =============================================================================
-# Header shims for CUDA 13 / NVTX removal
-# =============================================================================
-# (1) NVTX3: CUDA 13 dropped bundled NVTX headers; the pypi placeholder is empty.
-# Clone the upstream NVIDIA/NVTX repo and install just the C include tree.
-RUN git clone --depth 1 https://github.com/NVIDIA/NVTX.git /tmp/nvtx_src \
-    && cp -r /tmp/nvtx_src/c/include/nvtx3 /usr/local/cuda/include/nvtx3 \
-    && rm -rf /tmp/nvtx_src
-
-# (2) cuda_profiler_api.h: CUDA 13 removed the public header; libcudart.so.13
-# still exports the two symbols. TE 2.10 only #includes the header without
-# actually calling the APIs in the affected TUs, so a tiny shim is sufficient.
-COPY docker/patch/gb10/cuda_profiler_api.h /usr/local/cuda/include/cuda_profiler_api.h
-
-# =============================================================================
-# cuDNN engines_precompiled library
-# =============================================================================
-# NGC deliberately ships only runtime-compiled cuDNN libs to save space, but
-# TE's bundled cudnn-frontend requires CUDNN::cudnn_engines_precompiled. The
-# pypi wheel nvidia-cudnn-cu13 9.20.0.48 contains the missing .so.
-RUN PIP_CONSTRAINT= pip install --no-deps nvidia-cudnn-cu13==9.20.0.48 \
-    && ln -sf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/lib/libcudnn_engines_precompiled.so.9 \
-              /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so.9 \
-    && ln -sf /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so.9 \
-              /usr/lib/aarch64-linux-gnu/libcudnn_engines_precompiled.so
-
-# =============================================================================
-# CMake 4.3.1: needed for CUDA 13 family-specific `f` arch suffix (120f/121f).
-# NGC's /etc/pip/constraint.txt pins cmake==3.31.6; we must clear PIP_CONSTRAINT
-# for this install. CMAKE_POLICY_VERSION_MINIMUM=3.5 lets CMake 4.x accept
-# bundled third-party CMakeLists that declare cmake_minimum_required < 3.5.
-# =============================================================================
-RUN PIP_CONSTRAINT= pip install --no-deps --ignore-installed --force-reinstall cmake==4.3.1
-ENV CMAKE_POLICY_VERSION_MINIMUM=3.5
-
-# =============================================================================
-# Build prerequisites for the source installs below
-# =============================================================================
-RUN pip install --no-deps pybind11 pathspec pyproject_metadata scikit-build-core
-
-# =============================================================================
-# sgl-kernel for sm_121a (SGLang's CUDA kernel library)
-# =============================================================================
-# The published sgl_kernel wheels are cu12x (libnvrtc.so.12 ABI), plus the cu130
-# arm64 wheel only has the sm_100 variant, not sm_121. Also ABI-mismatch against
-# NGC torch. Build from source with the GB10 arch whitelist patch.
-#
-# Patch: docker/patch/gb10/patch_sgl_kernel.py adds SGL_KERNEL_GB10_ONLY CMake
-# option that restricts gencode to sm_120a + sm_121a (skipping sm_90a/sm_100a/
-# sm_103a/sm_110a). This is essential — compiling all seven archs simultaneously
-# OOM-kills cicc on 128 GB Spark hosts.
-ARG SGLANG_COMMIT=v0.5.9
-RUN git clone --depth 1 --branch ${SGLANG_COMMIT} https://github.com/sgl-project/sglang.git /root/src/sglang
-COPY docker/patch/gb10/patch_sgl_kernel.py /tmp/patch_sgl_kernel.py
-RUN python3 /tmp/patch_sgl_kernel.py /root/src/sglang/sgl-kernel/CMakeLists.txt \
-    && rm /tmp/patch_sgl_kernel.py
-
-# Conservative parallelism (8 parallel × 2 nvcc-internal threads = 16 active)
-# keeps peak memory ~25 GB on the 128 GB Spark. ~25 minutes on 20-core GB10.
-RUN cd /root/src/sglang/sgl-kernel \
-    && CMAKE_BUILD_PARALLEL_LEVEL=8 \
-       CMAKE_ARGS="-DSGL_KERNEL_COMPILE_THREADS=2 -DSGL_KERNEL_GB10_ONLY=ON" \
-       pip wheel . --no-build-isolation --no-deps -w /tmp/wheels \
-    && pip install --no-deps /tmp/wheels/sgl_kernel-*.whl \
-    && rm -rf /tmp/wheels /root/.cache/pip
-
-# =============================================================================
-# SGLang runtime + the small pure-python deps it needs.
-# --no-deps preserves NGC torch / transformers / flashinfer / quack-kernels /
-# cutlass-dsl / xgrammar / outlines that were carefully pinned in the NGC base.
-# Pulling sglang's full dep set here would downgrade torch to 2.9.1+cu129 and
-# break everything.
-# =============================================================================
-RUN pip install --no-deps \
-        "sglang==0.5.9" \
-        "torch_memory_saver==0.0.9" \
-        IPython traitlets stack_data executing asttokens pure_eval prompt_toolkit wcwidth \
-        orjson hf_transfer modelscope litellm
-
-# SGLang router: zhuzilin's fork ships x86_64-only. Upstream 0.3.2 has an arm64
-# wheel that satisfies slime's version-parse comparisons (`<=0.2.1`, `<0.3.0`).
-# Only loss: the `'slime' in version` check in wandb_utils.py won't match,
-# which is a non-critical logging branch.
-RUN pip install --no-deps sglang-router==0.3.2
-
-# =============================================================================
-# TransformerEngine 2.10 for sm_121f
-# =============================================================================
-# GB10 needs the CUDA 13 "f" (family-specific) arch suffix — TE's ptx.cuh has
-# static_asserts that reject plain sm_120 / sm_121 in favor of sm_120f / sm_121f.
-# CMake 4 is required to pass this through CMAKE_CUDA_ARCHITECTURES.
-RUN TORCH_CUDA_ARCH_LIST="12.0+PTX" \
-    NVTE_CUDA_ARCHS="120f;121f" \
-    NVTE_FRAMEWORK=pytorch \
-    MAX_JOBS=8 CMAKE_BUILD_PARALLEL_LEVEL=8 \
-    pip install --no-build-isolation --no-deps \
-        "git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.10"
-
-# TE pytorch subpackage pulls onnxscript → onnx_ir → onnx → ml_dtypes at import
-# time (not setup time). Install the runtime chain so `import transformer_engine.pytorch` works.
-RUN pip install --no-deps onnxscript onnx_ir onnx ml_dtypes
-
-# =============================================================================
-# apex (NVIDIA fused optimizers / norms)
-# =============================================================================
-# Pinned commit matches slime upstream Dockerfile.
-RUN NVCC_APPEND_FLAGS="--threads 4" MAX_JOBS=8 \
-    pip install --no-cache-dir --no-build-isolation --no-deps \
-        --config-settings "--build-option=--cpp_ext --cuda_ext --parallel 8" \
-        "git+https://github.com/NVIDIA/apex.git@10417aceddd7d5d05d7cbf7b0fc2daad1105f8b4"
-
-# =============================================================================
-# Megatron-LM (+ slime's megatron.patch)
-# =============================================================================
-# slime pins Megatron commit 3714d81d...
-ARG MEGATRON_COMMIT=3714d81d418c9f1bca4594fc35f9e8289f652862
-RUN git clone https://github.com/NVIDIA/Megatron-LM.git /root/src/Megatron-LM \
-    && cd /root/src/Megatron-LM \
-    && git checkout ${MEGATRON_COMMIT}
-COPY docker/patch/latest/megatron.patch /tmp/megatron.patch
-RUN cd /root/src/Megatron-LM \
-    && git apply /tmp/megatron.patch --3way \
-    && ! grep -R -n '^<<<<<<< ' . \
-    && rm /tmp/megatron.patch \
-    && pip install --no-deps --no-build-isolation -e .
-
-# Megatron.training / .rl / .legacy are sibling dirs meant to be on PYTHONPATH
-# (slime's own docs confirm this).
-ENV PYTHONPATH="/root/src/Megatron-LM"
-
-# =============================================================================
-# Remaining slime deps: bridges, modelopt, FLA, tilelang, antlr, omegaconf chain
-# =============================================================================
-RUN pip install --no-deps \
-        "git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c" \
-        flash-linear-attention==0.4.1 \
-        tilelang==0.1.8
-
-RUN pip install --no-deps --no-build-isolation \
-        "git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl" \
-        "nvidia-modelopt>=0.37.0"
-
-# antlr4-python3-runtime must be pinned to 4.9.3 — omegaconf's bundled grammar
-# was generated with antlr 4.9 serialized format and runtime 4.13 breaks on it
-# with "Could not deserialize ATN with version 3 (expected 4)".
-RUN pip install --no-deps "antlr4-python3-runtime==4.9.3" omegaconf wandb
-
-# Megatron requires numpy 1.x; NGC ships 2.x.
-RUN PIP_CONSTRAINT= pip install --no-deps --force-reinstall "numpy<2"
-
-# Reward function runtime deps (grade_answer_verl / deepscaler math verifier)
-RUN pip install --no-deps pylatexenc math_verify word2number
-
-# Eval / data-prep runtime deps
-RUN pip install --no-deps pyarrow accelerate
-
-# =============================================================================
-# slime itself
-# =============================================================================
-COPY . /root/slime
-RUN cd /root/slime \
-    && pip install --no-deps --no-build-isolation -e .
-
-# slime's int4 QAT kernel (fused per-channel fake-quant + dequant)
-RUN cd /root/slime/slime/backends/megatron_utils/kernels/int4_qat \
-    && pip install --no-build-isolation --no-deps .
-
-WORKDIR /root/slime
-
-# Sanity: load the full arg parser on container build
-RUN python train.py --help > /tmp/help.txt \
-    && test $(wc -l < /tmp/help.txt) -gt 3000 \
-    && echo "slime train.py --help OK ($(wc -l < /tmp/help.txt) lines)" \
-    && rm /tmp/help.txt
-
-CMD ["/bin/bash"]
diff --git a/docker/NOTES_GB10.md b/docker/NOTES_GB10.md
deleted file mode 100644
index d5417321f5..0000000000
--- a/docker/NOTES_GB10.md
+++ /dev/null
@@ -1,99 +0,0 @@
-# slime on NVIDIA DGX Spark (GB10) — Porting Notes
-
-Target platform:
-- Chip: NVIDIA GB10 (Grace + consumer Blackwell, SM 12.1 / sm_121a)
-- Arch: aarch64
-- OS: Ubuntu 24.04, NVIDIA driver 580.142 (CUDA 13.x forward-compat)
-- Unified memory: 128 GB (CPU+GPU shared)
-
-## Why a new Dockerfile is needed
-
-slime's published Docker images (`slimerl/slime:*`) are x86_64-only. The arm64 base
-it derives from (`slimerl/sglang:v0.5.9`) ships with CUDA 12.9 and stock PyTorch
-2.9.1+cu129 — neither of which knows about the `sm_121a` target used by GB10:
-the Triton PTX pipeline crashes with `ptxas fatal: Value 'sm_121a' is not defined`.
-
-The `ENABLE_CUDA_13=1` branch in the upstream Dockerfile is aimed at GB200/GB300
-(sm_100a) and still uses an x86-only router wheel and amd64 base, so it does not
-directly apply to GB10.
-
-This port rebases slime on `nvcr.io/nvidia/vllm:26.03-py3` (arm64), which ships:
-- CUDA 13.2 (ptxas understands sm_121a ✅)
-- PyTorch 2.11.0a0 compiled with `compute_120` (PTX forward-compat to sm_121 ✅)
-- Triton 3.6.0 (verified: Triton kernel JITs on GB10 ✅)
-- flash-attn 2.7.4.post1 preinstalled
-
-### Base image pinning (for reproducibility)
-
-Pull from NGC and verify digest:
-
-```bash
-docker pull nvcr.io/nvidia/vllm:26.03-py3
-docker inspect nvcr.io/nvidia/vllm:26.03-py3 --format '{{range .RepoDigests}}{{.}}{{end}}'
-# Expected digest:
-#   nvcr.io/nvidia/vllm@sha256:13e327dad79e6e417f6687fec2ba76b0386d597082ec0ee003c1e964ec6ad0e7
-```
-
-All downstream steps pin this digest. Product page:
-https://catalog.ngc.nvidia.com/orgs/nvidia/containers/vllm?version=26.03-py3
-
-## Known blockers and their resolutions
-
-| # | Blocker | Root cause | Resolution |
-|---|---------|-----------|------------|
-| 1 | `ptxas fatal: sm_121a` | CUDA 12.9 ptxas predates sm_121a (added in CUDA 13.0) | Rebase on NGC CUDA 13.2 image |
-| 2 | `libnvrtc.so.12` missing from sgl_kernel wheel | Published sgl_kernel wheel is cu12x | Use cu130 wheel OR build from source |
-| 3 | `sgl_kernel/sm100/...abi3.so: undefined symbol _ZN3c104cuda29c10_cuda_check_implementationEiPKcS2_ib` | Wheel built against stock libtorch, NGC libtorch has different C++ ABI | Build sgl_kernel from source against NGC torch |
-| 4 | sgl_kernel arm64 wheels only build sm_100 variant | sgl-project CI doesn't target GB10 | Build sm_121 variant from source |
-| 5 | sgl-kernel source build: cicc OOM on sm_90+100+103+110+120+121 concurrent compile | Cutlass FP8 templates × 7 arches × 12 parallel → >128GB | `SGL_KERNEL_GB10_ONLY=ON` CMake option (see `docker/patch/gb10/sgl-kernel-arch.patch`), drops to sm_120a+121a |
-| 6 | TE 2.10 build: `CUDNN::cudnn_engines_precompiled` target not found | NGC vLLM image ships only runtime-compiled cuDNN libs | Install `nvidia-cudnn-cu13==9.20.0.48` pypi wheel (has the precompiled engines), symlink into `/usr/lib/aarch64-linux-gnu/` |
-| 7 | TE 2.10 build: `nvtx3/nvToolsExt.h: No such file` | CUDA 13 dropped bundled NVTX headers; `nvidia-nvtx-cu13` pypi wheel is an empty 0.0.1 placeholder | Clone `NVIDIA/NVTX` github, copy `c/include/nvtx3/*` to `/usr/local/cuda/include/nvtx3/` |
-| 8 | TE 2.10 build: `ptx.cuh` static_assert "Compile for smXXXf instead of smXXX" | Blackwell family-specific TMA features require CUDA 13's new `f` (family-specific) arch suffix, not plain integer | Set `NVTE_CUDA_ARCHS="120f;121f"` — but CMake ≤3.31 rejects `f` suffix |
-| 9 | CMake 3.31 rejects `CMAKE_CUDA_ARCHITECTURES=120f;121f` | `f` suffix for CUDA 13 Blackwell family is only supported in CMake ≥4.0 | Upgrade to `cmake==4.3.1` via pip (must override NGC `/etc/pip/constraint.txt` with `PIP_CONSTRAINT=`) and set `CMAKE_POLICY_VERSION_MINIMUM=3.5` for old bundled deps |
-| 10 | TE 2.10 build: `cuda_profiler_api.h: No such file` | CUDA 13 removed the public header for `cudaProfilerStart/Stop`; the symbols still exist in `libcudart.so.13`. TE's 3 `fused_softmax` TUs `#include` the header but don't call the APIs | Install a 20-line shim header at `/usr/local/cuda/include/cuda_profiler_api.h` declaring the two functions extern. Stored as `docker/patch/gb10/cuda_profiler_api.h` |
-| 11 | slime `train.py --help`: `'tuple' object has no attribute 'strip'` | Typo in `slime/utils/arguments.py:1073`: `help=("string",)` (trailing comma → tuple) instead of `help=("string")` | Remove trailing comma — simple one-line slime fix, upstream-able |
-| 12 | `sglang_router` x86_64-only wheel from `zhuzilin/sgl-router` fork | slime Dockerfile pins `zhuzilin/sgl-router` release (no arm64 builds); slime's `'slime' in version` assertion is only in `wandb_utils.py` (non-critical path) | Install upstream `sglang-router==0.3.2` from PyPI (has arm64 wheel). Accept wandb path fallback |
-| 13 | `antlr4-python3-runtime==4.13.2` → `Could not deserialize ATN with version 3` | Omegaconf's bundled grammar was generated with antlr 4.9 serialized format; runtime 4.13 only reads format v4 | Pin `antlr4-python3-runtime==4.9.3` |
-| 14 | `megatron.training` not importable after `pip install -e Megatron-LM` | Megatron-LM's setup.py only packages `megatron-core`; `megatron.training`, `megatron.rl`, `megatron.legacy` are sibling dirs meant to be on `PYTHONPATH` | `export PYTHONPATH=/root/src/Megatron-LM:$PYTHONPATH` (slime docs confirm this) |
-| 15 | `libz3.so` missing for tilelang | tilelang uses Z3 SMT solver for autoscheduling; NGC vllm base doesn't include libz3 | `apt-get update && apt-get install -y libz3-dev` (libz3-4 package alias needs update first) |
-
-## Build journal
-
-### sgl-kernel OOM during build → arch whitelist patch
-
-First build attempt on GB10 with stock sgl-kernel v0.5.9 CMake: cicc
-(NVIDIA's CUDA frontend compiler) was OOM-killed while compiling
-`csrc/gemm/fp8_gemm_kernel.cu`, `fp8_blockwise_gemm_kernel.cu`, and
-`nvfp4_scaled_mm_kernels.cu`.
-
-Root cause: for `CUDA_VERSION >= 13.0 && aarch64`, sgl-kernel unconditionally
-emits seven gencodes per TU — `sm_90, sm_90a, sm_100a, sm_103a, sm_110a,
-sm_120a, sm_121a`. Cutlass FP8/NVFP4 gemm template instantiation uses
-~10–15 GB of RAM per TU per arch. Combined with 12 parallel nvcc jobs, peak
-memory exceeded the 128 GB unified memory limit.
-
-Only `sm_121a` (and `sm_120a` as PTX fallback) actually runs on GB10. Hopper
-(sm_90a) and datacenter Blackwell (sm_100a/sm_103a) binaries are dead weight.
-
-Fix: `docker/patch/gb10/sgl-kernel-arch.patch` adds a CMake option
-`SGL_KERNEL_GB10_ONLY`. When set, the other gencodes and FA3 (sm_90a-only)
-are skipped. Default OFF preserves upstream behavior. See the patch python
-script `docker/patch/gb10/patch_sgl_kernel.py` for a programmatic apply.
-
-### sgl-kernel build success (M1)
-
-With `SGL_KERNEL_GB10_ONLY=ON`, `CMAKE_BUILD_PARALLEL_LEVEL=8`, and
-`SGL_KERNEL_COMPILE_THREADS=2`, the build completed in 24 minutes with RAM
-never exceeding ~25 GB. Produced `sgl_kernel-0.3.21-cp310-abi3-linux_aarch64.whl`
-(74 MB), ABI-compatible with NGC libtorch 2.11.0a0 (verified via clean import
-and functional rmsnorm kernel on GB10).
-
-## Work items
-
-- [x] M1: Build sgl_kernel from source for sm_121 + NGC torch ABI
-- [x] M2: TransformerEngine 2.10 (sm_121f), apex, Megatron-LM all built + verified
-- [x] M3: slime installed with megatron.patch applied; `train.py --help` prints 3714-line arg list
-- [x] M4: End-to-end RL loop (Qwen2.5-0.5B + dapo-math-17k + GRPO colocated on 1 GB10 GPU) runs clean in 2m10s
-
-Reproducible image: `slime-gb10:m4-success` (36.3 GB, committed post-smoke)
-- Step 0 train metrics logged; full rollout → reward → policy-update → weight-sync cycle confirmed.
diff --git a/docker/README.md b/docker/README.md
index 62105e3ede..69b99fe7b9 100644
--- a/docker/README.md
+++ b/docker/README.md
@@ -5,10 +5,9 @@ We will publish 2 kinds of docker images:
 2. latest version, which aligns to `lmsysorg/sglang:latest`.
 
 current stable version is:
-- sglang v0.5.9 (bbe9c7eeb520b0a67e92d133dfc137a3688dc7f2), megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862
+- sglang v0.5.7 nightly-dev-20260107-dce8b060 (dce8b0606c06d3a191a24c7b8cbe8e238ab316c9), megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862
 
 history versions:
-- sglang v0.5.7 nightly-dev-20260107-dce8b060 (dce8b0606c06d3a191a24c7b8cbe8e238ab316c9), megatron dev 3714d81d418c9f1bca4594fc35f9e8289f652862
 - sglang v0.5.6 nightly-dev-20251208-5e2cda61 (5e2cda6158e670e64b926a9985d65826c537ac82), megatron v0.14.0 (23e00ed0963c35382dfe8a5a94fb3cda4d21e133)
 - sglang v0.5.5.post1 (303cc957e62384044dfa8e52d7d8af8abe12f0ac), megatron v0.14.0 (23e00ed0963c35382dfe8a5a94fb3cda4d21e133)
 - sglang v0.5.0rc0-cu126 (8ecf6b9d2480c3f600826c7d8fef6a16ed603c3f), megatron 48406695c4efcf1026a7ed70bb390793918dd97b
diff --git a/docker/npu_patch/README.md b/docker/npu_patch/README.md
deleted file mode 100644
index db93b67633..0000000000
--- a/docker/npu_patch/README.md
+++ /dev/null
@@ -1,192 +0,0 @@
-# Slime NPU Patch Installation Guide
-
-This guide provides instructions for installing Slime with NPU support, including all required dependencies and patches.
-
-## Component Version Mapping
-
-| Component       | Version/Commit                           | Source                                                                                                              |
-| --------------- | ---------------------------------------- | ------------------------------------------------------------------------------------------------------------------- |
-| Slime           | v0.2.2                                   | [GitHub](https://github.com/THUDM/slime/tree/v0.2.2)                                                                |
-| SGLang          | dce8b0606c06d3a191a24c7b8cbe8e238ab316c9 | [GitHub](https://github.com/sgl-project/sglang/tree/sglang-slime)                                             |
-| SGL Kernel NPU  | 2026.02.01                               | [GitHub](https://github.com/sgl-project/sgl-kernel-npu/releases/tag/2026.02.01)                                     |
-| Megatron-Bridge | 35b4ebfc486fb15dcc0273ceea804c3606be948a | [GitHub](https://github.com/fzyzcjy/Megatron-Bridge)                                                                |
-| Megatron-LM     | 3714d81d418c9f1bca4594fc35f9e8289f652862 | [GitHub](https://github.com/NVIDIA/Megatron-LM)                                                                     |
-| MindSpeed       | fc63de5c48426dd019c3b3f39e65f5bdf56e4086 | [GitCode](https://gitcode.com/Ascend/MindSpeed)                                                                     |
-| HDK             | 25.3.RC1                                 | [Ascend](https://www.hiascend.com/hardware/firmware-drivers/commercial?product=7\&model=33)                         |
-| CANN            | 8.5.0                                    | [Ascend](https://www.hiascend.com/developer/download/community/result?module=cann\&cann=8.5.0\&product=7\&model=33) |
-
-## Preparing the Running Environment
-
-### Python Version
-
-Only `python==3.11` is supported currently.
-
-```shell
-conda create -n slime_release python=3.11
-conda activate slime_release
-```
-
-### Working Directory Setup
-
-```shell
-mkdir <WORKSPACE> && cd <WORKSPACE>
-```
-
-### CANN Environment
-
-Prior to start work with Slime on Ascend you need to install CANN Toolkit, Kernels operator package and NNAL version 8.5.0, check the [installation guide](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1/softwareinst/instg/instg_0008.html?Mode=PmIns\&InstallType=local\&OS=openEuler\&Software=cannToolKit)
-
-```shell
-source <CANN_PATH>/ascend-toolkit/set_env.sh
-source <CANN_PATH>/nnal/atb/set_env.sh
-```
-
-### PyTorch and PyTorch NPU
-
-```shell
-pip install torch-npu==2.8.0
-```
-
-## Installing Dependencies
-
-### SGLang
-
-```shell
-cd <WORKSPACE>
-git clone https://github.com/sgl-project/sglang.git && cd sglang
-git checkout dce8b0606c06d3a191a24c7b8cbe8e238ab316c9
-mv python/pyproject.toml python/pyproject.toml.backup
-mv python/pyproject_other.toml python/pyproject.toml
-pip install -e "python[srt_npu]"
-pip install torch-npu==2.8.0
-```
-
-### SGL Kernel NPU and Torch Memory Saver
-
-Download `sgl-kernel-npu-2026.02.01-torch2.8.0-py311-cann8.5.0-a3-aarch64.zip` from the release link, then install:
-
-```shell
-pip install sgl_kernel_npu-2026.2.1-cp311-cp311-linux_aarch64.whl
-pip install torch_memory_saver-0.0.8-cp311-cp311-linux_aarch64.whl
-```
-
-### Megatron-Bridge
-
-```shell
-pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
-
-cd <WORKSPACE>
-git clone https://github.com/fzyzcjy/Megatron-Bridge.git -b dev_rl
-pip install nvidia-modelopt[torch]>=0.37.0 --no-build-isolation
-```
-
-### Megatron-LM
-
-```shell
-cd <WORKSPACE>
-git clone https://github.com/NVIDIA/Megatron-LM.git --recursive && \
-  cd Megatron-LM/ && git checkout 3714d81d418c9f1bca4594fc35f9e8289f652862 && \
-  pip install -e .
-```
-
-### MindSpeed
-
-```shell
-cd <WORKSPACE>
-git clone https://gitcode.com/Ascend/MindSpeed.git && \
-  cd MindSpeed/ && git checkout fc63de5c48426dd019c3b3f39e65f5bdf56e4086 && \
-  pip install -e .
-```
-
-### Slime
-
-```shell
-cd <WORKSPACE>
-git clone https://github.com/ascend-slime/slime.git && cd slime
-cp -r docker/npu_patch ../npu_patch
-git checkout v0.2.2
-pip install -e .
-```
-
-## Applying Patches
-
-```shell
-cd <WORKSPACE>/slime
-git apply ../npu_patch/slime.patch
-
-cd <WORKSPACE>/sglang
-git apply ../slime/docker/patch/v0.5.7/sglang.patch
-git apply ../npu_patch/sglang.patch
-
-cd <WORKSPACE>/Megatron-LM
-git apply ../slime/docker/patch/v0.5.7/megatron.patch
-git apply ../npu_patch/megatron.patch
-
-cd <WORKSPACE>/Megatron-Bridge
-git apply ../npu_patch/megatron-bridge.patch
-
-cd <WORKSPACE>/MindSpeed
-git apply ../npu_patch/mindspeed.patch
-```
-
-## Additional Dependencies
-
-```shell
-cd <WORKSPACE>/slime
-pip install triton-ascend
-pip install torch-npu==2.8.0
-pip install torchvision==0.23.0
-pip install numpy==1.26.0
-```
-
-## Running the Training
-
-### Configuration
-
-Modify the paths in the following files according to your environment (note to use your CANN version):
-
-**Common (both GRPO and PPO):**
-
-- `slime/utils/external_utils/command_utils.py`
-
-**GRPO:**
-
-- `examples/geo3k_vlm_multi_turn/run_grpo_npu.sh`
-- `examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py`
-
-**PPO:**
-
-- `examples/geo3k_vlm_multi_turn/run_ppo_npu.sh`
-- `examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py`
-
-### Dataset
-
-Download the dataset from [HuggingFace](https://huggingface.co/datasets/VeraIsHere/geo3k_imgurl_processed) following the instructions in the script directory.
-
-### Execute Training
-
-```shell
-cd <WORKSPACE>/slime
-# GRPO
-bash examples/geo3k_vlm_multi_turn/run_grpo_npu.sh
-# PPO
-bash examples/geo3k_vlm_multi_turn/run_ppo_npu.sh
-```
-
-To save logs and display them simultaneously:
-
-```shell
-# GRPO
-bash examples/geo3k_vlm_multi_turn/run_grpo_npu.sh 2>&1 | tee -a <LOG_FILE>
-# PPO
-bash examples/geo3k_vlm_multi_turn/run_ppo_npu.sh 2>&1 | tee -a <LOG_FILE>
-```
-
-## Placeholders Reference
-
-| Placeholder   | Description                          | Example               |
-| ------------- | ------------------------------------ | --------------------- |
-| `<WORKSPACE>` | Root directory for all installations | `/root/slime-release` |
-| `<CANN_PATH>` | Path to CANN installation directory  | `/usr/local/ascend`   |
-| `<LOG_FILE>`  | Path to log file for training output | `training.log`        |
-
diff --git a/docker/npu_patch/megatron-bridge.patch b/docker/npu_patch/megatron-bridge.patch
deleted file mode 100644
index 3817097f40..0000000000
--- a/docker/npu_patch/megatron-bridge.patch
+++ /dev/null
@@ -1,93 +0,0 @@
-diff --git a/src/megatron/bridge/models/conversion/param_mapping.py b/src/megatron/bridge/models/conversion/param_mapping.py
-index dc7d0be..8156826 100644
---- a/src/megatron/bridge/models/conversion/param_mapping.py
-+++ b/src/megatron/bridge/models/conversion/param_mapping.py
-@@ -1088,15 +1088,19 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
-             "ColumnParallelLinear",
-             "TEColumnParallelLinear",
-             "TELayerNormColumnParallelLinear",
-+            "MindSpeedTELayerNormColumnParallelLinear",
-             "TEColumnParallelGroupedLinear",
-+            "MindSpeedTEColumnParallelGroupedLinear",
-             "VocabParallelEmbedding",
-             "DotProductAttention",  # for attention sink only
-             "TEDotProductAttention",  # for attention sink only
-+            "MindSpeedTEDotProductAttention",
-         },
-         "row": {
-             "RowParallelLinear",
-             "TERowParallelLinear",
-             "TERowParallelGroupedLinear",
-+            "MindSpeedTERowParallelGroupedLinear",
-         },
-         "replicated": {
-             # Normalization layers
-@@ -1164,7 +1168,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
-         # Handle fused modules like TELayerNormColumnParallelLinear
-         # These modules have both column-parallel weights (weight, bias)
-         # and replicated layer norm weights (layer_norm_weight, layer_norm_bias)
--        if module_type == "TELayerNormColumnParallelLinear":
-+        if module_type == "TELayerNormColumnParallelLinear" or module_type == "MindSpeedTELayerNormColumnParallelLinear":
-             # Check the actual parameter name to determine the correct parallelism type
-             if self.megatron_param and (
-                 self.megatron_param.endswith("layer_norm_weight") or self.megatron_param.endswith("layer_norm_bias")
-@@ -1195,7 +1199,7 @@ class AutoMapping(MegatronParamMapping[torch.Tensor]):
-             return "replicated"
- 
-         # Check parallel_mode for TELinear
--        if module_type == "TELinear":
-+        if module_type == "TELinear" or module_type == "MindSpeedTELinear":
-             if module.parallel_mode == "column":
-                 return "column"
-             elif module.parallel_mode == "row":
-diff --git a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
-index 3f64d8d..9dadc4b 100644
---- a/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
-+++ b/src/megatron/bridge/models/qwen_vl/modelling_qwen3_vl/transformer_block.py
-@@ -71,8 +71,9 @@ class Qwen3VLTransformerBlock(TransformerBlock):
-                 context_mask,
-                 rotary_pos_emb,
-                 visual_pos_masks,
--                deepstack_visual_embeds,
-+                *deepstack_visual_embeds_args,
-             ):
-+                deepstack_visual_embeds = list(deepstack_visual_embeds_args) if deepstack_visual_embeds_args else None
-                 for index in range(start, end):
-                     layer = self._get_layer(index)
-                     inner_fp8_context = (
-@@ -103,6 +104,8 @@ class Qwen3VLTransformerBlock(TransformerBlock):
-                 return hidden_states, context
- 
-             return custom_forward
-+        
-+        deepstack_visual_embeds_tuple = tuple(deepstack_visual_embeds) if deepstack_visual_embeds else ()
- 
-         def checkpoint_handler(forward_func):
-             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
-@@ -118,7 +121,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
-                     context_mask,
-                     rotary_pos_emb,
-                     visual_pos_masks,
--                    deepstack_visual_embeds,
-+                    *deepstack_visual_embeds_tuple,
-                 )
-             else:
-                 return tensor_parallel.checkpoint(
-@@ -130,7 +133,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
-                     context_mask,
-                     rotary_pos_emb,
-                     visual_pos_masks,
--                    deepstack_visual_embeds,
-+                    *deepstack_visual_embeds_tuple,
-                 )
- 
-         if self.config.recompute_method == "uniform":
-@@ -169,7 +172,7 @@ class Qwen3VLTransformerBlock(TransformerBlock):
-                         context_mask,
-                         rotary_pos_emb,
-                         visual_pos_masks,
--                        deepstack_visual_embeds,
-+                        *deepstack_visual_embeds_tuple,
-                     )
-         else:
-             raise ValueError("Invalid activation recompute method.")
diff --git a/docker/npu_patch/megatron.patch b/docker/npu_patch/megatron.patch
deleted file mode 100644
index 7687827db6..0000000000
--- a/docker/npu_patch/megatron.patch
+++ /dev/null
@@ -1,518 +0,0 @@
-diff --git a/megatron/core/activations.py b/megatron/core/activations.py
-index 8b422d73a..58fba4667 100644
---- a/megatron/core/activations.py
-+++ b/megatron/core/activations.py
-@@ -5,19 +5,19 @@ import torch.nn.functional as F
- from megatron.core.jit import jit_fuser
- 
- 
--@jit_fuser
-+
- def squared_relu(x: torch.Tensor) -> torch.Tensor:
-     """Squared ReLU activation"""
-     return torch.pow(F.relu(x), 2)
- 
- 
--@jit_fuser
-+
- def quick_gelu(x: torch.Tensor) -> torch.Tensor:
-     """Quick GELU activation"""
-     return x * torch.sigmoid(1.702 * x)
- 
- 
--@jit_fuser
-+
- def fast_gelu(x: torch.Tensor) -> torch.Tensor:
-     """Fast GELU activation"""
-     return 0.5 * x * (1.0 + torch.tanh(x * 0.7978845608 * (1.0 + 0.044715 * x * x)))
-diff --git a/megatron/core/fusions/fused_bias_dropout.py b/megatron/core/fusions/fused_bias_dropout.py
-index 336452562..614ee1a48 100644
---- a/megatron/core/fusions/fused_bias_dropout.py
-+++ b/megatron/core/fusions/fused_bias_dropout.py
-@@ -64,14 +64,14 @@ def bias_dropout_add_unfused(training):
-     return _bias_dropout_add
- 
- 
--@jit_fuser
-+
- def bias_dropout_add_fused_train(
-     x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
- ) -> torch.Tensor:
-     return _bias_dropout_add_func(x_with_bias, residual, prob, True)
- 
- 
--@jit_fuser
-+
- def bias_dropout_add_fused_inference(
-     x_with_bias: Tuple[torch.Tensor, Optional[torch.Tensor]], residual: torch.Tensor, prob: float
- ) -> torch.Tensor:
-diff --git a/megatron/core/fusions/fused_bias_geglu.py b/megatron/core/fusions/fused_bias_geglu.py
-index 7a7fbe7f9..f9a438953 100644
---- a/megatron/core/fusions/fused_bias_geglu.py
-+++ b/megatron/core/fusions/fused_bias_geglu.py
-@@ -13,7 +13,7 @@ from megatron.core.jit import jit_fuser
- # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
- 
- 
--@jit_fuser
-+
- def geglu(y):
-     """Performs GEGLU (GELU-Gated Linear Unit) activation.
- 
-@@ -27,7 +27,7 @@ def geglu(y):
-     return (y_1 * 0.5 * (1.0 + torch.tanh(0.79788456 * y_1 * (1 + 0.044715 * y_1 * y_1)))) * y_2
- 
- 
--@jit_fuser
-+
- def bias_geglu(bias, y):
-     """Performs GEGLU activation with bias addition.
- 
-@@ -45,7 +45,7 @@ def bias_geglu(bias, y):
- # gradient of tanh approximation of gelu
- # gradient of actual gelu is:
- # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
--@jit_fuser
-+
- def geglu_back(g, y):
-     """Computes the gradient for the GEGLU activation.
- 
-@@ -65,7 +65,7 @@ def geglu_back(g, y):
-     return torch.cat(((g * y_2) * ff, g * (y_1 * 0.5 * (1.0 + tanh_out))), -1)
- 
- 
--@jit_fuser
-+
- def bias_geglu_back(g, y, bias):
-     """Computes the gradient for the biased GEGLU activation.
- 
-@@ -181,13 +181,13 @@ def bias_geglu_impl(input, bias):
- # ------------------------- QUICK GEGLU FUSION --------------------------
- 
- 
--@jit_fuser
-+
- def quick_gelu(y: torch.Tensor) -> torch.Tensor:
-     """Sigmoid approximation of gelu"""
-     return y * torch.sigmoid(1.702 * y)
- 
- 
--@jit_fuser
-+
- def quick_geglu(y: torch.Tensor, linear_offset: float = 0.0) -> torch.Tensor:
-     """Performs Quick-GELU-based GEGLU activation : quick_gelu(y1) * (y2 + offset).
- 
-@@ -202,7 +202,7 @@ def quick_geglu(y: torch.Tensor, linear_offset: float = 0.0) -> torch.Tensor:
-     return quick_gelu(y_1) * (y_2 + linear_offset)
- 
- 
--@jit_fuser
-+
- def weighted_quick_geglu(
-     y: torch.Tensor, weights: torch.Tensor, linear_offset: float = 0.0
- ) -> torch.Tensor:
-@@ -217,7 +217,7 @@ def weighted_quick_geglu(
- 
- 
- # gradient of sigmoid approximation of gelu
--@jit_fuser
-+
- def quick_geglu_back(g, y, linear_offset: float = 0.0) -> torch.Tensor:
-     """Backward helper for Quick-GEGLU.
- 
-@@ -236,7 +236,7 @@ def quick_geglu_back(g, y, linear_offset: float = 0.0) -> torch.Tensor:
-     return torch.cat((dy_1, dy_2), -1)
- 
- 
--@jit_fuser
-+
- def weighted_quick_geglu_back(g, y, weights, linear_offset: float = 0.0):
-     """Backward helper for weighted Quick-GEGLU.
-     Returns gradient w.r.t input `y` and `weights`.
-@@ -255,7 +255,7 @@ def weighted_quick_geglu_back(g, y, weights, linear_offset: float = 0.0):
- # ---------------- Weighted Bias Quick-GEGLU helpers -----------------
- 
- 
--@jit_fuser
-+
- def weighted_bias_quick_geglu(
-     y: torch.Tensor, bias: torch.Tensor, weights: torch.Tensor, linear_offset: float = 0.0
- ) -> torch.Tensor:
-@@ -275,7 +275,7 @@ def weighted_bias_quick_geglu(
-     return res.to(dtype)
- 
- 
--@jit_fuser
-+
- def weighted_bias_quick_geglu_back(g, y, bias, weights, linear_offset: float = 0.0):
-     """Backward helper for weighted Quick-GEGLU with bias.
- 
-diff --git a/megatron/core/fusions/fused_bias_gelu.py b/megatron/core/fusions/fused_bias_gelu.py
-index 8cc90f617..fda8f2f5f 100644
---- a/megatron/core/fusions/fused_bias_gelu.py
-+++ b/megatron/core/fusions/fused_bias_gelu.py
-@@ -13,7 +13,7 @@ from megatron.core.jit import jit_fuser
- # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
- 
- 
--@jit_fuser
-+
- def bias_gelu(bias, y):
-     x = bias + y
-     return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
-@@ -22,7 +22,7 @@ def bias_gelu(bias, y):
- # gradient of tanh approximation of gelu
- # gradient of actual gelu is:
- # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
--@jit_fuser
-+
- def bias_gelu_back(g, bias, y):
-     x = bias + y
-     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
-diff --git a/megatron/core/fusions/fused_bias_swiglu.py b/megatron/core/fusions/fused_bias_swiglu.py
-index 632470876..105936786 100644
---- a/megatron/core/fusions/fused_bias_swiglu.py
-+++ b/megatron/core/fusions/fused_bias_swiglu.py
-@@ -12,7 +12,7 @@ from megatron.core.utils import nvtx_decorator
- ###### BIAS SWIGLU FUSION/ NO AUTOGRAD ################
- 
- 
--@jit_fuser
-+
- def swiglu(y):
-     """Performs SwiGLU (Swish-Gated Linear Unit) activation function.
- 
-@@ -26,7 +26,7 @@ def swiglu(y):
-     return F.silu(y_1) * y_2
- 
- 
--@jit_fuser
-+
- def bias_swiglu(y, bias):
-     """Performs SwiGLU activation with bias addition.
- 
-@@ -41,7 +41,7 @@ def bias_swiglu(y, bias):
-     return swiglu(y)
- 
- 
--@jit_fuser
-+
- def weighted_swiglu(y, weights):
-     dtype = y.dtype
-     res = swiglu(y) * weights
-@@ -51,7 +51,7 @@ def weighted_swiglu(y, weights):
- # gradient of tanh approximation of gelu
- # gradient of actual gelu is:
- # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
--@jit_fuser
-+
- def swiglu_back(g, y):
-     """Computes the gradient for the SwiGLU activation function.
- 
-@@ -69,7 +69,7 @@ def swiglu_back(g, y):
-     )
- 
- 
--@jit_fuser
-+
- def bias_swiglu_back(g, y, bias):
-     """Computes the gradient for the biased SwiGLU activation function.
- 
-@@ -86,7 +86,7 @@ def bias_swiglu_back(g, y, bias):
-     return swiglu_back(g, y)
- 
- 
--@jit_fuser
-+
- def weighted_swiglu_back(g, y, weights):
-     input_dtype = y.dtype
-     w_dtype = weights.dtype
-diff --git a/megatron/core/fusions/fused_cross_entropy.py b/megatron/core/fusions/fused_cross_entropy.py
-index 23e4b6031..4d447e0e0 100644
---- a/megatron/core/fusions/fused_cross_entropy.py
-+++ b/megatron/core/fusions/fused_cross_entropy.py
-@@ -9,7 +9,7 @@ from megatron.core.tensor_parallel.cross_entropy import VocabParallelCrossEntrop
- from megatron.core.tensor_parallel.utils import VocabUtility
- 
- 
--@jit_fuser
-+
- def calculate_logits_max(vocab_parallel_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
-     """
-     Calculates the maximum logits of the predicted tokens.
-@@ -22,7 +22,7 @@ def calculate_logits_max(vocab_parallel_logits: torch.Tensor) -> Tuple[torch.Ten
-     return vocab_parallel_logits, logits_max
- 
- 
--@jit_fuser
-+
- def calculate_predicted_logits(
-     vocab_parallel_logits: torch.Tensor,
-     target: torch.Tensor,
-@@ -44,7 +44,7 @@ def calculate_predicted_logits(
-     return target_mask, masked_target_1d, predicted_logits_sum_exp_logits, exp_logits
- 
- 
--@jit_fuser
-+
- def calculate_cross_entropy_loss(
-     exp_logits: torch.Tensor, predicted_logits_sum_exp_logits: torch.Tensor
- ) -> Tuple[torch.Tensor, torch.Tensor]:
-@@ -61,7 +61,7 @@ def calculate_cross_entropy_loss(
-     return exp_logits, loss
- 
- 
--@jit_fuser
-+
- def calculate_gradients(
-     softmax: torch.Tensor,
-     grad_output: torch.Tensor,
-diff --git a/megatron/core/fusions/fused_pad_routing_map.py b/megatron/core/fusions/fused_pad_routing_map.py
-index c382178b6..563279edd 100644
---- a/megatron/core/fusions/fused_pad_routing_map.py
-+++ b/megatron/core/fusions/fused_pad_routing_map.py
-@@ -70,7 +70,7 @@ def _pad_routing_map_kernel(
-     tl.store(output_row_ptr + token_indices, output_row, mask=token_mask)
- 
- 
--@jit_fuser
-+
- def fused_pad_routing_map(routing_map: torch.Tensor, pad_multiple: int) -> torch.Tensor:
-     """Fused version of pad_routing_map.
-     Args:
-diff --git a/megatron/core/fusions/fused_weighted_squared_relu.py b/megatron/core/fusions/fused_weighted_squared_relu.py
-index 02dabc14c..137022386 100644
---- a/megatron/core/fusions/fused_weighted_squared_relu.py
-+++ b/megatron/core/fusions/fused_weighted_squared_relu.py
-@@ -10,7 +10,7 @@ from megatron.core.utils import nvtx_decorator
- ######################  WEIGHTED SQUARED ReLU FUSION  ######################
- 
- 
--@jit_fuser
-+
- def weighted_squared_relu(x: torch.Tensor, weights: torch.Tensor) -> torch.Tensor:
-     """Element-wise weight applied after Squared-ReLU.
- 
-@@ -28,7 +28,7 @@ def weighted_squared_relu(x: torch.Tensor, weights: torch.Tensor) -> torch.Tenso
-     return res.to(out_dtype)
- 
- 
--@jit_fuser
-+
- def _squared_relu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-     """Gradient of Squared-ReLU.
- 
-@@ -37,7 +37,7 @@ def _squared_relu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
-     return g * 2 * F.relu(x)
- 
- 
--@jit_fuser
-+
- def weighted_squared_relu_back(g: torch.Tensor, x: torch.Tensor, weights: torch.Tensor):
-     """Backward for weighted Squared-ReLU.
- 
-diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
-index 712793853..76ea4333b 100755
---- a/megatron/core/models/gpt/gpt_layer_specs.py
-+++ b/megatron/core/models/gpt/gpt_layer_specs.py
-@@ -219,7 +219,7 @@ def get_gpt_layer_with_transformer_engine_spec(
-             'The fp8 argument in "get_gpt_layer_with_transformer_engine_spec" has been deprecated'
-             " and will be removed soon. Please update your code accordingly."
-         )
--
-+    from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider
-     if use_kitchen:
-         assert HAVE_KITCHEN
-         backend: BackendSpecProvider = KitchenSpecProvider(
-diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py
-index dfa6e4c35..8b3621819 100644
---- a/megatron/core/ssm/gated_delta_net.py
-+++ b/megatron/core/ssm/gated_delta_net.py
-@@ -413,7 +413,7 @@ class GatedDeltaNet(MegatronModule):
- 
-         return out, out_bias
- 
--    @jit_fuser
-+    
-     def _apply_gated_norm(self, x, gate):
-         # Output Norm
-         x_dtype = x.dtype
-diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
-index 80e9ec6fc..b6bcef6a9 100644
---- a/megatron/core/transformer/attention.py
-+++ b/megatron/core/transformer/attention.py
-@@ -1026,7 +1026,7 @@ class Attention(MegatronModule, ABC):
- 
-         return output, bias
- 
--    @jit_fuser
-+    
-     def _apply_output_gate(self, x, gate):
-         x_dtype = x.dtype
-         gate = gate.contiguous()
-diff --git a/megatron/core/transformer/module.py b/megatron/core/transformer/module.py
-index 2330df91b..0446c9097 100644
---- a/megatron/core/transformer/module.py
-+++ b/megatron/core/transformer/module.py
-@@ -16,9 +16,9 @@ from megatron.core.transformer.utils import (
-     sharded_state_dict_default,
- )
- 
--_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
--_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
--_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)
-+_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor, torch.npu.FloatTensor)
-+_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor, torch.npu.HalfTensor)
-+_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor, torch.npu.BFloat16Tensor)
- 
- 
- def param_is_not_shared(param):  # pylint: disable=missing-function-docstring
-diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py
-index 5eeafdd8d..a4dce6970 100644
---- a/megatron/core/transformer/moe/experts.py
-+++ b/megatron/core/transformer/moe/experts.py
-@@ -91,7 +91,7 @@ class GroupedMLP(MegatronModule):
-             if self.config.activation_func not in (F.silu, F.gelu):
-                 raise ValueError("Activation function must be silu or gelu when using GroupedMLP.")
- 
--            @jit_fuser
-+            
-             def glu(x):
-                 x = torch.chunk(x, 2, dim=-1)
-                 return self.config.activation_func(x[0]) * x[1]
-@@ -108,7 +108,7 @@ class GroupedMLP(MegatronModule):
-                 "moe_act recompute for fp8 or fp4 cannot work with the legacy GroupedMLP."
-             )
- 
--        @jit_fuser
-+        
-         def activation_func_with_probs(x, probs):
-             dtype = x.dtype
-             res = self.activation_func(x) * probs
-diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
-index 517944f25..2c5bc0728 100644
---- a/megatron/core/transformer/moe/router.py
-+++ b/megatron/core/transformer/moe/router.py
-@@ -472,7 +472,7 @@ class TopKRouter(Router):
-         else:
-             return input
- 
--    @jit_fuser
-+    
-     def _apply_expert_bias(self, routing_map: torch.Tensor):
-         """
-         Update expert bias and tokens_per_expert
-diff --git a/megatron/core/transformer/moe/token_dispatcher.py b/megatron/core/transformer/moe/token_dispatcher.py
-index d0da38d63..6d092c88f 100644
---- a/megatron/core/transformer/moe/token_dispatcher.py
-+++ b/megatron/core/transformer/moe/token_dispatcher.py
-@@ -1403,7 +1403,7 @@ class MoEFlexTokenDispatcher(MoETokenDispatcher):
-         ).contiguous()
-         return routing_map, probs
- 
--    @jit_fuser
-+    
-     def dispatch_preprocess(
-         self, hidden_states: torch.Tensor, routing_map: torch.Tensor, probs: torch.Tensor
-     ):
-diff --git a/megatron/core/transformer/torch_norm.py b/megatron/core/transformer/torch_norm.py
-index d0ceca7af..f16796680 100644
---- a/megatron/core/transformer/torch_norm.py
-+++ b/megatron/core/transformer/torch_norm.py
-@@ -69,7 +69,7 @@ class L2Norm(torch.nn.Module):
-         self.hidden_size = hidden_size
-         self.eps = eps
- 
--    @jit_fuser
-+    
-     def _norm(self, x):
-         """
-         Performs the actual L2 normalization.
-diff --git a/megatron/core/transformer/utils.py b/megatron/core/transformer/utils.py
-index 880c53099..dbc95736e 100644
---- a/megatron/core/transformer/utils.py
-+++ b/megatron/core/transformer/utils.py
-@@ -51,7 +51,7 @@ def attention_mask_func(attention_scores, attention_mask):
-     return attention_scores
- 
- 
--@jit_fuser
-+
- def gelu_impl(x):
-     """OpenAI's gelu implementation."""
-     return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
-@@ -65,7 +65,7 @@ def openai_gelu(x):
- # This is actually Python equivalent of torch.nn.functional.gelu(), also with
- # type hints for ONNX exporter
- # pylint: disable=missing-function-docstring
--@jit_fuser
-+
- def erf_gelu(x):
-     return (
-         x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype))
-diff --git a/megatron/legacy/model/fused_bias_gelu.py b/megatron/legacy/model/fused_bias_gelu.py
-index e00e63148..ffe4b7ec6 100644
---- a/megatron/legacy/model/fused_bias_gelu.py
-+++ b/megatron/legacy/model/fused_bias_gelu.py
-@@ -12,7 +12,7 @@ from megatron.core.jit import jit_fuser
- # actual gelu is:
- # x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
- 
--@jit_fuser
-+
- def bias_gelu(bias, y):
-     x = bias + y
-     return  x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
-@@ -20,7 +20,7 @@ def bias_gelu(bias, y):
- # gradient of tanh approximation of gelu
- # gradient of actual gelu is:
- # 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
--@jit_fuser
-+
- def bias_gelu_back(g, bias, y):
-     x = bias + y
-     tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
-diff --git a/megatron/legacy/model/transformer.py b/megatron/legacy/model/transformer.py
-index 2a662a55b..2fc3e1bfb 100644
---- a/megatron/legacy/model/transformer.py
-+++ b/megatron/legacy/model/transformer.py
-@@ -856,7 +856,7 @@ def get_bias_dropout_add(training):
-     return _bias_dropout_add
- 
- 
--@jit_fuser
-+
- def bias_dropout_add_fused_train(x: torch.Tensor,
-                                  bias: Optional[torch.Tensor],
-                                  residual: torch.Tensor,
-@@ -864,7 +864,7 @@ def bias_dropout_add_fused_train(x: torch.Tensor,
-     return bias_dropout_add(x, bias, residual, prob, True)
- 
- 
--@jit_fuser
-+
- def bias_dropout_add_fused_inference(x: torch.Tensor,
-                                      bias: Optional[torch.Tensor],
-                                      residual: torch.Tensor,
-diff --git a/megatron/legacy/model/utils.py b/megatron/legacy/model/utils.py
-index 5762000d5..534858df7 100644
---- a/megatron/legacy/model/utils.py
-+++ b/megatron/legacy/model/utils.py
-@@ -43,7 +43,7 @@ def get_linear_layer(rows, columns, init_method):
-     return layer
- 
- 
--@jit_fuser
-+
- def gelu_impl(x):
-     """OpenAI's gelu implementation."""
-     return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x *
-@@ -54,7 +54,7 @@ def openai_gelu(x):
- 
- 
- #This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter
--@jit_fuser
-+
- def erf_gelu(x):
-     return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype)+torch.ones_like(x).to(dtype=x.dtype))
- 
diff --git a/docker/npu_patch/mindspeed.patch b/docker/npu_patch/mindspeed.patch
deleted file mode 100644
index 0cf1c84840..0000000000
--- a/docker/npu_patch/mindspeed.patch
+++ /dev/null
@@ -1,48 +0,0 @@
-diff --git a/mindspeed/core/fusions/fused_rope.py b/mindspeed/core/fusions/fused_rope.py
-index a6f02e07..70f7cb08 100644
---- a/mindspeed/core/fusions/fused_rope.py
-+++ b/mindspeed/core/fusions/fused_rope.py
-@@ -126,5 +126,6 @@ def apply_rotary_pos_emb(
-         freqs,
-         rotary_interleaved=config.rotary_interleaved,
-         multi_latent_attention=config.multi_latent_attention,
--        mscale=mscale
-+        mscale=mscale,
-+        cp_group=cp_group
-     )
-diff --git a/mindspeed/megatron_adaptor.py b/mindspeed/megatron_adaptor.py
-index f14b231d..4615590e 100644
---- a/mindspeed/megatron_adaptor.py
-+++ b/mindspeed/megatron_adaptor.py
-@@ -57,6 +57,7 @@ def delete_lock_file():
- def repatch(args):
-     MindSpeedFeaturesManager.remove_patches()
-     full_args = get_full_args()
-+    args = vars(args)
-     for k, v in args.items():
-         setattr(full_args, k, v)
-     MindSpeedFeaturesManager.apply_features_pre_patches(full_args)
-diff --git a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
-index ac4eabe5..78c5866d 100644
---- a/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
-+++ b/mindspeed/te/pytorch/attention/dot_product_attention/dot_product_attention.py
-@@ -330,6 +330,8 @@ class DotProductAttention(torch.nn.Module):
-         inference_params: Any = None,
-         pad_between_seqs: Optional[bool] = None,
-         fp8_output: Optional[bool] = False,
-+        local_cp_size=None,
-+        cp_group=None,
-     ) -> torch.Tensor:
-         """
-         Dot Product Attention Layer.
-@@ -659,7 +661,9 @@ class MindSpeedTEDotProductAttention(DotProductAttention):
-         ) and not getattr(self.config, 'is_llava', False):
-             self.config.sparse_mode = 2
-             attention_mask = get_attention_mask(self.config)
--
-+        attention_mask = torch.triu(
-+            torch.ones((2048, 2048),
-+                       device=query.device, dtype=torch.bool), diagonal=1)
-         packed_seq_kwargs = (
-             {key: getattr(packed_seq_params, key) for key in self.kept_packed_seq_params}
-             if packed_seq_params is not None
diff --git a/docker/npu_patch/sglang.patch b/docker/npu_patch/sglang.patch
deleted file mode 100644
index b684d3c0eb..0000000000
--- a/docker/npu_patch/sglang.patch
+++ /dev/null
@@ -1,610 +0,0 @@
-diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
-index 8a343d43b..3e49dddb0 100644
---- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
-+++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
-@@ -1050,93 +1050,36 @@ class AscendAttnBackend(AttentionBackend):
-                 )
- 
-         if not self.use_mla:
--            num_tokens = q.shape[0]
--            """PA will support bs<tp in the later version of CANN"""
--            if num_tokens < get_attention_tp_size():
--                k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
--                    layer.layer_id
--                ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
--                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
--                    layer.layer_id
--                ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
--                query = q.reshape(-1, 1, layer.tp_q_head_num * layer.qk_head_dim)
--                if self.forward_metadata.seq_lens_cpu_int is None:
--                    actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
--                else:
--                    actual_seq_len_kv = (
--                        self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
--                    )
--                num_tokens = query.shape[0]
--                workspace = (
--                    torch_npu._npu_fused_infer_attention_score_get_max_workspace(
--                        query,
--                        k_cache,
--                        v_cache,
--                        block_table=self.forward_metadata.block_tables,
--                        block_size=self.page_size,
--                        num_heads=layer.tp_q_head_num,
--                        num_key_value_heads=layer.tp_k_head_num,
--                        input_layout="BSH",
--                        scale=layer.scaling,
--                        actual_seq_lengths_kv=actual_seq_len_kv,
--                    )
--                )
--                output = torch.empty(
--                    (num_tokens, 1, layer.tp_q_head_num * layer.v_head_dim),
--                    dtype=q.dtype,
--                    device=q.device,
--                )
--                softmax_lse = torch.empty(1, dtype=q.dtype, device=q.device)
--                torch_npu.npu_fused_infer_attention_score.out(
--                    query,
--                    k_cache,
--                    v_cache,
--                    block_table=self.forward_metadata.block_tables,
--                    block_size=self.page_size,
--                    num_heads=layer.tp_q_head_num,
--                    num_key_value_heads=layer.tp_k_head_num,
--                    input_layout="BSH",
--                    scale=layer.scaling,
--                    actual_seq_lengths_kv=actual_seq_len_kv,
--                    workspace=workspace,
--                    out=[output, softmax_lse],
--                )
--                return output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
-+            k_cache = forward_batch.token_to_kv_pool.get_key_buffer(
-+                layer.layer_id
-+            ).view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim)
-+            v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
-+                layer.layer_id
-+            ).view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim)
-+            query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
-+            num_tokens = query.shape[0]
-+            if self.forward_metadata.seq_lens_cpu_int is None:
-+                actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
-             else:
--                k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id)
--                v_cache = forward_batch.token_to_kv_pool.get_value_buffer(
--                    layer.layer_id
--                )
--                query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim)
--                num_tokens = query.shape[0]
--                attn_output = torch.empty(
--                    (num_tokens, layer.tp_q_head_num, layer.v_head_dim),
--                    dtype=query.dtype,
--                    device=query.device,
--                )
--                if self.forward_metadata.seq_lens_cpu_int is None:
--                    actual_seq_len_kv = torch.from_numpy(
--                        np.array(self.forward_metadata.seq_lens_cpu_list).astype(
--                            np.int32
--                        )
--                    )
--                else:
--                    actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_int
--
--                torch_npu._npu_paged_attention(
--                    query=query,
--                    key_cache=k_cache,
--                    value_cache=v_cache,
--                    num_heads=layer.tp_q_head_num,
--                    num_kv_heads=layer.tp_k_head_num,
--                    scale_value=layer.scaling,
--                    block_table=self.forward_metadata.block_tables,
--                    context_lens=actual_seq_len_kv,
--                    out=attn_output,
--                )
--                return attn_output.view(
--                    num_tokens, layer.tp_q_head_num * layer.v_head_dim
-+                actual_seq_len_kv = (
-+                    self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist()
-                 )
-+            actual_seq_lengths = np.arange(1, num_tokens + 1)
-+            attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
-+                query,
-+                k_cache,
-+                v_cache,
-+                block_table=self.forward_metadata.block_tables,
-+                block_size=self.page_size,
-+                num_heads=layer.tp_q_head_num,
-+                num_key_value_heads=layer.tp_k_head_num,
-+                input_layout="TND",
-+                scale=layer.scaling,
-+                actual_seq_lengths=actual_seq_lengths,
-+                actual_seq_lengths_kv=actual_seq_len_kv,
-+                sparse_mode=0,
-+            )
-+            return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim)
-         else:
-             c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
-             if is_fia_nz():
-diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
-index d089fcc63..ce19d2cd0 100644
---- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
-+++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
-@@ -29,6 +29,7 @@ import sglang
- from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
- from sglang.srt.distributed.parallel_state import GroupCoordinator
- from sglang.srt.layers.dp_attention import get_attention_tp_size
-+from sglang.srt.environ import envs
- from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
- from sglang.srt.utils import (
-     empty_context,
-@@ -111,23 +112,15 @@ class NPUGraphRunner(CudaGraphRunner):
-             out = run_once_fn()
-         return out
- 
--    def _get_update_attr_name(self, model_runner, forward_batch):
--        if (
--            self.bs < get_attention_tp_size()
--            or forward_batch.forward_mode.is_target_verify()
--            or self.use_fia
--        ):
--            return self.attr_name[AttentionArch.MLA]
--        return self.attr_name[model_runner.model_config.attention_arch]
--
--    def _get_update_attr_type(self, model_runner, forward_batch):
--        if (
--            self.bs < get_attention_tp_size()
--            or forward_batch.forward_mode.is_target_verify()
--            or self.use_fia
--        ):
--            return self.attr_type[AttentionArch.MLA]
--        return self.attr_type[model_runner.model_config.attention_arch]
-+    def _get_update_attr_name(self):
-+        # if envs.SGLANG_USE_PAGED_ATTENTION.get():
-+        #     return self.attr_name[AttentionArch.MHA]
-+        return self.attr_name[AttentionArch.MLA]
-+
-+    def _get_update_attr_type(self):
-+        # if envs.SGLANG_USE_PAGED_ATTENTION.get():
-+        #     return self.attr_type[AttentionArch.MHA]
-+        return self.attr_type[AttentionArch.MLA]
- 
-     def _update_inputs(self, seq_lens):
-         if isinstance(self.update_attr_type, torch.Tensor):
-@@ -144,6 +137,10 @@ class NPUGraphRunner(CudaGraphRunner):
-         output_dir = os.path.join(
-             os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile"
-         )
-+        if ".." in output_dir:
-+            raise ValueError(
-+                "Output directory path contains '..', which is not allowed."
-+            )
-         if not Path(output_dir).exists():
-             Path(output_dir).mkdir(parents=True, exist_ok=True)
-         logger.info(
-@@ -181,12 +178,8 @@ class NPUGraphRunner(CudaGraphRunner):
-             self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
-             self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
- 
--        self.update_attr_name = self._get_update_attr_name(
--            self.model_runner, forward_batch
--        )
--        self.update_attr_type = self._get_update_attr_type(
--            self.model_runner, forward_batch
--        )
-+        self.update_attr_name = self._get_update_attr_name()
-+        self.update_attr_type = self._get_update_attr_type()
-         # Replay
-         if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
-             if forward_batch.forward_mode.is_target_verify():
-diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py
-index dd8ca7d4f..92be78315 100644
---- a/python/sglang/srt/layers/rotary_embedding.py
-+++ b/python/sglang/srt/layers/rotary_embedding.py
-@@ -115,9 +115,10 @@ class RotaryEmbedding(MultiPlatformOp):
-             cache = cache.to(dtype)
- 
-         if (
--            (not (_is_cuda or _is_npu) or self.head_size not in [64, 128, 256, 512])
-+            (not (_is_cuda) or self.head_size not in [64, 128, 256, 512])
-             and not (_is_cpu and _is_cpu_amx_available)
-             and not (_is_xpu)
-+            and not (_is_npu)
-         ):
-             if _is_cuda or _is_hip:
-                 from sgl_kernel import rotary_embedding
-@@ -292,8 +293,8 @@ class RotaryEmbedding(MultiPlatformOp):
-         assert (
-             fused_set_kv_buffer_arg is None
-         ), "fused_set_kv_buffer_arg is not supported for npu implementation"
--
--        rotary_mode = "half"
-+        if query.dtype == torch.bfloat16 and self.cos_sin_cache.dtype == torch.float:
-+            return self.forward_native(positions, query, key, offsets)
-         if self.is_neox_style:
-             rotary_mode = "half"
-         else:
-diff --git a/python/sglang/srt/layers/vocab_parallel_embedding.py b/python/sglang/srt/layers/vocab_parallel_embedding.py
-index f171a15f2..192970904 100644
---- a/python/sglang/srt/layers/vocab_parallel_embedding.py
-+++ b/python/sglang/srt/layers/vocab_parallel_embedding.py
-@@ -31,6 +31,7 @@ from sglang.srt.utils import (
-     cpu_has_amx_support,
-     get_compiler_backend,
-     is_cpu,
-+    is_npu,
-     set_weight_attrs,
- )
- 
-@@ -38,6 +39,7 @@ DEFAULT_VOCAB_PADDING_SIZE = 64
- 
- _is_cpu_amx_available = cpu_has_amx_support()
- _is_cpu = is_cpu()
-+_is_npu = is_npu()
- 
- logger = logging.getLogger(__name__)
- 
-@@ -123,7 +125,7 @@ class VocabParallelEmbeddingShardIndices:
-         assert self.num_added_elements <= self.num_added_elements_padded
- 
- 
--@torch.compile(dynamic=True, backend=get_compiler_backend())
-+@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
- def get_masked_input_and_mask(
-     input_: torch.Tensor,
-     org_vocab_start_index: int,
-diff --git a/python/sglang/srt/models/llama.py b/python/sglang/srt/models/llama.py
-index 4176af28d..ffcf73d46 100644
---- a/python/sglang/srt/models/llama.py
-+++ b/python/sglang/srt/models/llama.py
-@@ -52,10 +52,14 @@ from sglang.srt.model_loader.weight_utils import (
-     maybe_remap_kv_scale_name,
- )
- from sglang.srt.server_args import get_global_server_args
--from sglang.srt.utils import add_prefix, make_layers
-+from sglang.srt.utils import add_prefix, is_npu, make_layers
- from sglang.utils import get_exception_traceback
- 
- logger = logging.getLogger(__name__)
-+_is_npu = is_npu()
-+
-+if _is_npu:
-+    from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope
- 
- 
- class LlamaMLP(nn.Module):
-@@ -185,15 +189,44 @@ class LlamaAttention(nn.Module):
-             prefix=add_prefix("attn", prefix),
-         )
- 
-+    def forward_prepare_native(self, positions, hidden_states):
-+        qkv, _ = self.qkv_proj(hidden_states)
-+        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
-+        q, k = self.rotary_emb(positions, q, k)
-+        return q, k, v
-+
-+    def forward_prepare_npu(self, positions, hidden_states, forward_batch):
-+        qkv, _ = self.qkv_proj(hidden_states)
-+        if self.attn.layer_id == forward_batch.token_to_kv_pool.start_layer:
-+            self.rotary_emb.get_cos_sin_with_position(positions)
-+        q, k, v = split_qkv_rmsnorm_rope(
-+            qkv,
-+            self.rotary_emb.position_sin,
-+            self.rotary_emb.position_cos,
-+            self.q_size,
-+            self.kv_size,
-+            self.head_dim,
-+        )
-+        return q, k, v
-+
-     def forward(
-         self,
-         positions: torch.Tensor,
-         hidden_states: torch.Tensor,
-         forward_batch: ForwardBatch,
-     ) -> torch.Tensor:
--        qkv, _ = self.qkv_proj(hidden_states)
--        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
--        q, k = self.rotary_emb(positions, q, k)
-+        if not _is_npu or not hasattr(self.rotary_emb, "get_cos_sin_with_position"):
-+            q, k, v = self.forward_prepare_native(
-+                positions=positions,
-+                hidden_states=hidden_states,
-+            )
-+        else:
-+            q, k, v = self.forward_prepare_npu(
-+                positions=positions,
-+                hidden_states=hidden_states,
-+                forward_batch=forward_batch,
-+            )
-+
-         attn_output = self.attn(q, k, v, forward_batch)
-         output, _ = self.o_proj(attn_output)
-         return output
-diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py
-index 47a1a4e4c..2cdbf8906 100644
---- a/python/sglang/srt/models/qwen3.py
-+++ b/python/sglang/srt/models/qwen3.py
-@@ -161,12 +161,12 @@ class Qwen3Attention(nn.Module):
-             qkv,
-             self.rotary_emb.position_sin,
-             self.rotary_emb.position_cos,
--            self.q_norm.weight,
--            self.k_norm.weight,
-             self.q_size,
-             self.kv_size,
-             self.head_dim,
--            self.q_norm.variance_epsilon,
-+            eps=self.q_norm.variance_epsilon,
-+            q_weight=self.q_norm.weight,
-+            k_weight=self.k_norm.weight,
-             q_bias=getattr(self.q_norm, "bias", None),
-             k_bias=getattr(self.k_norm, "bias", None),
-         )
-@@ -370,6 +370,7 @@ class Qwen3ForCausalLM(nn.Module):
-                     config.vocab_size,
-                     config.hidden_size,
-                     quant_config=quant_config,
-+                    use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
-                     prefix=add_prefix("lm_head", prefix),
-                 )
-         else:
-diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py
-index e277d46f2..9e49156da 100644
---- a/python/sglang/srt/models/qwen3_moe.py
-+++ b/python/sglang/srt/models/qwen3_moe.py
-@@ -549,12 +549,12 @@ class Qwen3MoeAttention(nn.Module):
-             qkv,
-             self.rotary_emb.position_sin,
-             self.rotary_emb.position_cos,
--            self.q_norm.weight,
--            self.k_norm.weight,
-             self.q_size,
-             self.kv_size,
-             self.head_dim,
--            self.q_norm.variance_epsilon,
-+            eps=self.q_norm.variance_epsilon,
-+            q_weight=self.q_norm.weight,
-+            k_weight=self.k_norm.weight,
-             q_bias=getattr(self.q_norm, "bias", None),
-             k_bias=getattr(self.k_norm, "bias", None),
-         )
-diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py
-index 218e32362..9f3141ce2 100644
---- a/python/sglang/srt/models/qwen3_vl.py
-+++ b/python/sglang/srt/models/qwen3_vl.py
-@@ -19,6 +19,7 @@ import re
- from functools import lru_cache, partial
- from typing import Callable, Iterable, List, Optional, Tuple, Union
- 
-+import numpy as np
- import torch
- import torch.nn as nn
- from einops import rearrange
-@@ -397,70 +398,89 @@ class Qwen3VLMoeVisionModel(nn.Module, RotaryPosMixin):
-         return cos_combined, sin_combined
- 
-     def fast_pos_embed_interpolate(self, grid_thw):
--        grid_ts, grid_hs, grid_ws = grid_thw[:, 0], grid_thw[:, 1], grid_thw[:, 2]
-         num_grid_per_side = int(self.num_position_embeddings**0.5)
--        device = self.pos_embed.weight.device
- 
-         idx_list = [[] for _ in range(4)]
-         weight_list = [[] for _ in range(4)]
- 
--        for t, h, w in zip(grid_ts, grid_hs, grid_ws):
--            h_idxs = torch.linspace(0, num_grid_per_side - 1, h)
--            w_idxs = torch.linspace(0, num_grid_per_side - 1, w)
-+        # TODO: use torch instand of np
-+        for t, h, w in grid_thw:
-+            h_idxs = np.linspace(0, num_grid_per_side - 1, h)
-+            w_idxs = np.linspace(0, num_grid_per_side - 1, w)
- 
--            h_idxs_floor = h_idxs.int()
--            w_idxs_floor = w_idxs.int()
--            h_idxs_ceil = (h_idxs.int() + 1).clip(max=num_grid_per_side - 1)
--            w_idxs_ceil = (w_idxs.int() + 1).clip(max=num_grid_per_side - 1)
-+            h_idxs_floor = h_idxs.astype(int)
-+            w_idxs_floor = w_idxs.astype(int)
-+            h_idxs_ceil = (h_idxs.astype(int) + 1).clip(max=num_grid_per_side - 1)
-+            w_idxs_ceil = (w_idxs.astype(int) + 1).clip(max=num_grid_per_side - 1)
- 
-             dh = h_idxs - h_idxs_floor
-             dw = w_idxs - w_idxs_floor
- 
--            base_h = h_idxs_floor * num_grid_per_side
--            base_h_ceil = h_idxs_ceil * num_grid_per_side
--
--            indices = [
--                (base_h[None].T + w_idxs_floor[None]).flatten(),
--                (base_h[None].T + w_idxs_ceil[None]).flatten(),
--                (base_h_ceil[None].T + w_idxs_floor[None]).flatten(),
--                (base_h_ceil[None].T + w_idxs_ceil[None]).flatten(),
--            ]
-+            idx_list[0].extend(
-+                ((h_idxs_floor * num_grid_per_side)[None].T + w_idxs_floor[None])
-+                .flatten()
-+                .tolist()
-+                * t
-+            )
-+            idx_list[1].extend(
-+                ((h_idxs_floor * num_grid_per_side)[None].T + w_idxs_ceil[None])
-+                .flatten()
-+                .tolist()
-+                * t
-+            )
-+            idx_list[2].extend(
-+                ((h_idxs_ceil * num_grid_per_side)[None].T + w_idxs_floor[None])
-+                .flatten()
-+                .tolist()
-+                * t
-+            )
-+            idx_list[3].extend(
-+                ((h_idxs_ceil * num_grid_per_side)[None].T + w_idxs_ceil[None])
-+                .flatten()
-+                .tolist()
-+                * t
-+            )
- 
--            weights = [
--                ((1 - dh)[None].T * (1 - dw)[None]).flatten(),
--                ((1 - dh)[None].T * dw[None]).flatten(),
--                (dh[None].T * (1 - dw)[None]).flatten(),
--                (dh[None].T * dw[None]).flatten(),
--            ]
-+            weight_list[0].extend(
-+                ((1 - dh)[None].T * (1 - dw)[None]).flatten().tolist() * t
-+            )
-+            weight_list[1].extend(((1 - dh)[None].T * dw[None]).flatten().tolist() * t)
-+            weight_list[2].extend((dh[None].T * (1 - dw)[None]).flatten().tolist() * t)
-+            weight_list[3].extend((dh[None].T * dw[None]).flatten().tolist() * t)
- 
--            for i in range(4):
--                idx_list[i].extend(indices[i].tolist())
--                weight_list[i].extend(weights[i].tolist())
-+        device = self.pos_embed.weight.device
-+        dtype = self.pos_embed.weight.dtype
- 
--        idx_tensor = torch.tensor(idx_list, dtype=torch.long, device=device)
--        weight_tensor = torch.tensor(
--            weight_list, dtype=self.pos_embed.weight.dtype, device=device
-+        p0 = (
-+            self.pos_embed(torch.tensor(idx_list[0], dtype=torch.long, device=device))
-+            * torch.tensor(weight_list[0], dtype=dtype, device=device)[:, None]
-         )
--        pos_embeds = self.pos_embed(idx_tensor).to(device) * weight_tensor[:, :, None]
--        patch_pos_embeds = pos_embeds[0] + pos_embeds[1] + pos_embeds[2] + pos_embeds[3]
--
--        patch_pos_embeds = patch_pos_embeds.split(
--            [h * w for h, w in zip(grid_hs, grid_ws)]
-+        p1 = (
-+            self.pos_embed(torch.tensor(idx_list[1], dtype=torch.long, device=device))
-+            * torch.tensor(weight_list[1], dtype=dtype, device=device)[:, None]
-+        )
-+        p2 = (
-+            self.pos_embed(torch.tensor(idx_list[2], dtype=torch.long, device=device))
-+            * torch.tensor(weight_list[2], dtype=dtype, device=device)[:, None]
-+        )
-+        p3 = (
-+            self.pos_embed(torch.tensor(idx_list[3], dtype=torch.long, device=device))
-+            * torch.tensor(weight_list[3], dtype=dtype, device=device)[:, None]
-         )
- 
-+        patch_pos_embeds = p0 + p1 + p2 + p3
-+        patch_pos_embeds = patch_pos_embeds.split([t * h * w for t, h, w in grid_thw])
-         patch_pos_embeds_permute = []
--        merge_size = self.spatial_merge_size
--        for pos_embed, t, h, w in zip(patch_pos_embeds, grid_ts, grid_hs, grid_ws):
--            pos_embed = pos_embed.repeat(t, 1)
-+        m_size = self.spatial_merge_size
-+        for pos_embed, (t, h, w) in zip(patch_pos_embeds, grid_thw):
-             pos_embed = (
--                pos_embed.view(
--                    t, h // merge_size, merge_size, w // merge_size, merge_size, -1
--                )
-+                pos_embed.view(t, h // m_size, m_size, w // m_size, m_size, -1)
-                 .permute(0, 1, 3, 2, 4, 5)
-                 .flatten(0, 4)
-             )
-             patch_pos_embeds_permute.append(pos_embed)
--        return torch.cat(patch_pos_embeds_permute)
-+        patch_pos_embeds = torch.cat(patch_pos_embeds_permute)
-+        return patch_pos_embeds
- 
-     def forward(
-         self,
-@@ -650,31 +670,20 @@ class Qwen3LLMModel(Qwen3Model):
-                     hidden_states + residual if residual is not None else hidden_states
-                 )
- 
--            deepstack_embeds = None
--            if input_deepstack_embeds is not None:
--                prev_layer_idx = layer_idx - 1
--                if prev_layer_idx in self.deepstack_embed_to_decoder_layer:
--                    sep = self.hidden_size * prev_layer_idx
--                    deepstack_embeds = input_deepstack_embeds[
--                        :, sep : sep + self.hidden_size
--                    ]
--
--            # SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
--            # See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
--            # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
--            # The order matters because addition with different tensors is not associative in practice.
-             hidden_states, residual = layer(
-                 positions,
-                 hidden_states,
-                 forward_batch,
-                 residual,
--                post_residual_addition=deepstack_embeds,
-             )
-+            # process deepstack
-+            if (
-+                input_deepstack_embeds is not None
-+                and layer_idx in self.deepstack_embed_to_decoder_layer
-+            ):
- 
--        # Handle deepstack for the last processed layer if it exists.
--        last_deepstack = self.get_deepstack_embeds(
--            self.end_layer - 1, input_deepstack_embeds
--        )
-+                sep = self.hidden_size * layer_idx
-+                hidden_states += input_deepstack_embeds[:, sep : sep + self.hidden_size]
- 
-         if not self.pp_group.is_last_rank:
-             return PPProxyTensors(
-@@ -688,9 +697,7 @@ class Qwen3LLMModel(Qwen3Model):
-                 if residual is None:
-                     hidden_states = self.norm(hidden_states)
-                 else:
--                    hidden_states, _ = self.norm(
--                        hidden_states, residual, post_residual_addition=last_deepstack
--                    )
-+                    hidden_states, _ = self.norm(hidden_states, residual)
- 
-         if len(aux_hidden_states) == 0:
-             return hidden_states
-@@ -805,15 +812,15 @@ class Qwen3VLForConditionalGeneration(nn.Module):
-         max_images_per_call = get_int_env_var("SGLANG_VLM_MAX_IMAGES_PER_VIT", 0)
- 
-         if max_patches_per_call == 0 and max_images_per_call == 0:
--            if self.use_data_parallel:
--                return run_dp_sharded_mrope_vision_model(
--                    self.visual,
--                    pixel_values,
--                    image_grid_thw.tolist(),
--                    rope_type="rope_3d",
--                )
--            else:
--                return self.visual(pixel_values, grid_thw=image_grid_thw)
-+            # if self.use_data_parallel:
-+            #     return run_dp_sharded_mrope_vision_model(
-+            #         self.visual,
-+            #         pixel_values,
-+            #         image_grid_thw.tolist(),
-+            #         rope_type="rope_3d",
-+            #     )
-+            # else:
-+            return self.visual(pixel_values, grid_thw=image_grid_thw)
- 
-         # compute the number of patches per image and the slice positions in pixel_values
-         grid_thw_list = (
-@@ -995,7 +1002,7 @@ class Qwen3VLForConditionalGeneration(nn.Module):
-                 name = name.replace(r"model.language_model.", r"model.")
-             layer_id = get_layer_id(name)
- 
--            if self.pp_group.is_last_rank and "model.embed_tokens.weight" in name:
-+            if self.pp_group.is_last_rank and "model.embed_tokens.weight" in name and self.config.tie_word_embeddings:
-                 if "lm_head.weight" in params_dict:
-                     lm_head_param = params_dict["lm_head.weight"]
-                     weight_loader = getattr(
-diff --git a/scripts/ci/npu_ci_install_dependency.sh b/scripts/ci/npu_ci_install_dependency.sh
-index 6172db3b4..8cc8fbf39 100755
---- a/scripts/ci/npu_ci_install_dependency.sh
-+++ b/scripts/ci/npu_ci_install_dependency.sh
-@@ -49,7 +49,7 @@ wget -O "${BISHENG_NAME}" "${BISHENG_URL}" && chmod a+x "${BISHENG_NAME}" && "./
- 
- 
- ### Install sgl-kernel-npu
--SGL_KERNEL_NPU_TAG="20251206"
-+SGL_KERNEL_NPU_TAG="2025.12.31"
- git clone --depth 1 https://github.com/sgl-project/sgl-kernel-npu.git --branch ${SGL_KERNEL_NPU_TAG}
- (cd sgl-kernel-npu && bash ./build.sh && ${PIP_INSTALL} output/deep_ep*.whl output/sgl_kernel_npu*.whl && cd "$(python3 -m pip show deep-ep | grep -E '^Location:' | awk '{print $2}')" && ln -s deep_ep/deep_ep_cpp*.so)
- 
diff --git a/docker/npu_patch/slime.patch b/docker/npu_patch/slime.patch
deleted file mode 100644
index 55e7c2d44e..0000000000
--- a/docker/npu_patch/slime.patch
+++ /dev/null
@@ -1,941 +0,0 @@
-diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py
-index a4666fbe..2a07086a 100644
---- a/slime/backends/megatron_utils/__init__.py
-+++ b/slime/backends/megatron_utils/__init__.py
-@@ -21,21 +21,35 @@ except ImportError:
-     logging.warning("deep_ep is not installed, some functionalities may be limited.")
- 
- try:
--    from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import (
--        Qwen3VLMoETextRotaryEmbedding,
--        Qwen3VLTextRotaryEmbedding,
--    )
--
--    def patch_rotary_embedding(cls):
--        _original_forward = cls.forward
--
--        def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
--            return _original_forward(self, *args, **kwargs)
-+    from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.text_model import Qwen3VLTextRotaryEmbedding, Qwen3VLMoETextRotaryEmbedding
-+        
-+    _original_forward = Qwen3VLTextRotaryEmbedding.forward
-+    _original_forward_1 = Qwen3VLMoETextRotaryEmbedding.forward
-+    
-+    def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
-+        return _original_forward(self, *args, **kwargs)
-+    def _patched_forward_1(self, *args, packed_seq_params=None, **kwargs):
-+        return _original_forward_1(self, *args, **kwargs)
-+    Qwen3VLTextRotaryEmbedding.forward = _patched_forward
-+    Qwen3VLMoETextRotaryEmbedding.forward = _patched_forward_1
-+except ImportError:
-+    pass
- 
--        cls.forward = _patched_forward
-+try:
-+    from mbridge.models.qwen3_vl.model import Qwen3VLModel
-+    _original_forward2 = Qwen3VLModel.forward
-+    def _patched_forward2(self, *args, loss_mask=None, **kwargs):
-+        return _original_forward2(self, *args, **kwargs)
-+    Qwen3VLModel.forward = _patched_forward2
-+except ImportError:
-+    pass
- 
--    patch_rotary_embedding(Qwen3VLTextRotaryEmbedding)
--    patch_rotary_embedding(Qwen3VLMoETextRotaryEmbedding)
-+try:
-+    from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel
-+    _original_forward3 = Qwen3VLModel.forward
-+    def _patched_forward3(self, *args, loss_mask=None, **kwargs):
-+        return _original_forward3(self, *args, **kwargs)
-+    Qwen3VLModel.forward = _patched_forward3
- except ImportError:
-     pass
- 
-diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py
-index 7bc4f910..a0b1f63a 100644
---- a/slime/backends/megatron_utils/actor.py
-+++ b/slime/backends/megatron_utils/actor.py
-@@ -8,6 +8,10 @@ from contextlib import nullcontext
- import ray
- import torch
- import torch.distributed as dist
-+from slime.utils.common import is_npu
-+if is_npu():
-+    import mindspeed.megatron_adaptor
-+    from mindspeed.megatron_adaptor import repatch
- from megatron.core import mpu
- from ray.actor import ActorHandle
- from torch_memory_saver import torch_memory_saver
-@@ -55,6 +59,8 @@ class MegatronTrainRayActor(TrainRayActor):
-         super().init(args, role, with_ref)
- 
-         init(args)
-+        if is_npu():
-+            repatch(args)
- 
-         if is_megatron_main_rank():
-             init_tracking(args, primary=False)
-@@ -596,8 +602,12 @@ class MegatronTrainRayActor(TrainRayActor):
- 
-         group_name = "actor_critic"
-         world_size = 2
-+        if is_npu():
-+            backend = "hccl"
-+        else:
-+            backend = "nccl"
-         self._actor_critic_groups = init_process_group(
--            backend="nccl",
-+            backend=backend,
-             init_method=f"tcp://{master_address}:{master_port}",
-             world_size=world_size,
-             rank=0 if self.role == "actor" else 1,
-diff --git a/slime/backends/megatron_utils/megatron_to_hf/__init__.py b/slime/backends/megatron_utils/megatron_to_hf/__init__.py
-index 84ff899a..a0e2e43c 100644
---- a/slime/backends/megatron_utils/megatron_to_hf/__init__.py
-+++ b/slime/backends/megatron_utils/megatron_to_hf/__init__.py
-@@ -7,7 +7,7 @@ from .processors import quantize_params, remove_padding
- from .qwen2 import convert_qwen2_to_hf
- from .qwen3_next import convert_qwen3_next_to_hf
- from .qwen3moe import convert_qwen3moe_to_hf
--
-+from .qwen3_vl import convert_qwen3vl_to_hf
- 
- # TODO unify w/ `convert_to_hf`
- def postprocess_hf_param(args, megatron_param_name, hf_param_name, param):
-@@ -40,7 +40,10 @@ def _convert_to_hf_core(args, model_name, name, param):
-     elif "qwen3next" in model_name:
-         converted_named_tensors = convert_qwen3_next_to_hf(args, name, param)
-     elif "qwen2" in model_name or "qwen3" in model_name:
--        converted_named_tensors = convert_qwen2_to_hf(args, name, param)
-+        if "qwen3vl" in model_name:
-+            converted_named_tensors = convert_qwen3vl_to_hf(args, name, param)
-+        else:
-+            converted_named_tensors = convert_qwen2_to_hf(args, name, param)
-     elif "deepseekv3" in model_name:
-         converted_named_tensors = convert_deepseekv3_to_hf(args, name, param)
- 
-diff --git a/slime/backends/megatron_utils/megatron_to_hf/qwen3_vl.py b/slime/backends/megatron_utils/megatron_to_hf/qwen3_vl.py
-new file mode 100644
-index 00000000..5a3aa56d
---- /dev/null
-+++ b/slime/backends/megatron_utils/megatron_to_hf/qwen3_vl.py
-@@ -0,0 +1,133 @@
-+import re
-+import torch
-+from megatron.core import parallel_state as mpu
-+
-+
-+def convert_qwen3vl_to_hf(args, name, param):
-+    """
-+    Convert Megatron-style Qwen3-VL parameter names to HF-style names.
-+    Supports both language model and vision model parameters.
-+
-+    Args:
-+        args: megatron model args (num_attention_heads, kv_channels, etc.)
-+        name: str, Megatron parameter name
-+        param: torch.Tensor, parameter value
-+
-+    Returns:
-+        List of tuples [(hf_name, hf_param), ...]
-+    """
-+
-+    hf_name_param = None
-+
-+    # ----------------------------
-+    # 1. language model & vision model parameters
-+    # ----------------------------
-+
-+    try:
-+        head_dim = args.kv_channels if args.kv_channels is not None else args.hidden_size // args.num_attention_heads
-+    except:
-+        head_dim = args.hidden_size // args.num_attention_heads
-+    value_num_per_group = args.num_attention_heads // args.num_query_groups
-+    language_num_layers = args.num_layers
-+    
-+    pp_size = args.pipeline_model_parallel_size
-+    pp_rank = mpu.get_pipeline_model_parallel_rank()
-+    assert language_num_layers % pp_size == 0
-+
-+    num_layers_per_rank = language_num_layers // pp_size
-+    offsets = pp_rank * num_layers_per_rank
-+
-+
-+
-+    # ----------------------------
-+    # 2. LM Embeddings & output
-+    # ----------------------------
-+    if name == "module.module.language_model.embedding.word_embeddings.weight":
-+        hf_name_param = [("model.language_model.embed_tokens.weight", param)]
-+    elif name == "module.module.language_model.decoder.final_layernorm.weight":
-+        hf_name_param = [("model.language_model.norm.weight", param)]
-+    elif name == "module.module.language_model.output_layer.weight":
-+        if not args.untie_embeddings_and_output_weights:
-+            return [("model.language_model.embed_tokens.weight", param)]
-+        else:
-+            return [("lm_head.weight", param)]
-+
-+    else:
-+
-+        decoder_layers_pattern = r"module\.module\.language_model.decoder\.layers\.(\d+)\.(.+)"
-+        vision_pattern = r"module\.module\.vision_model\.(.+)"
-+
-+        if match := re.match(decoder_layers_pattern, name):
-+            # ----------------------------
-+            # 3. Attention and MLP layers in language model
-+            # ----------------------------
-+
-+            layer_idx, rest = match.groups()
-+            layer_idx = str(int(layer_idx) + offsets)
-+            # Self-attention projection
-+            if rest == "self_attention.linear_proj.weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.self_attn.o_proj.weight", param)]
-+            
-+            elif rest == "self_attention.linear_qkv.weight":
-+                param = param.view(args.num_query_groups, -1, head_dim, args.hidden_size)
-+                q_param, k_param, v_param = torch.split(
-+                    param, split_size_or_sections=[value_num_per_group, 1, 1], dim=1
-+                )
-+                q_param = q_param.reshape(-1, args.hidden_size)
-+                k_param = k_param.reshape(-1, args.hidden_size)
-+                v_param = v_param.reshape(-1, args.hidden_size)
-+                hf_name_param = [
-+                    (f"model.language_model.layers.{layer_idx}.self_attn.q_proj.weight", q_param),
-+                    (f"model.language_model.layers.{layer_idx}.self_attn.k_proj.weight", k_param),
-+                    (f"model.language_model.layers.{layer_idx}.self_attn.v_proj.weight", v_param),
-+                ]
-+
-+            # MLP layers
-+            elif rest == "mlp.linear_fc1.weight":
-+                gate_weight, up_weight = param.chunk(2, dim=0)
-+
-+                hf_name_param = [
-+                    (f"model.language_model.layers.{layer_idx}.mlp.gate_proj.weight", gate_weight),
-+                    (f"model.language_model.layers.{layer_idx}.mlp.up_proj.weight", up_weight),
-+                ]
-+
-+            elif rest == "mlp.linear_fc2.weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.mlp.down_proj.weight", param)]
-+
-+            # LayerNorms
-+            elif rest == "self_attention.linear_qkv.layer_norm_weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.input_layernorm.weight", param)]
-+            elif rest == "mlp.linear_fc1.layer_norm_weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.post_attention_layernorm.weight", param)]
-+            elif rest == "self_attention.q_layernorm.weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.self_attn.q_norm.weight", param)]
-+            elif rest == "self_attention.k_layernorm.weight":
-+                hf_name_param = [(f"model.language_model.layers.{layer_idx}.self_attn.k_norm.weight", param)]
-+
-+        elif match_v := re.match(vision_pattern, name):
-+            # ----------------------------
-+            # 4. Vision model parameters
-+            # ----------------------------
-+            deepstack_merger_pattern = r"deepstack_merger_list\.(\d+)\.(.+)"
-+            decoder_layer_pattern = r"blocks\.(\d+)\.(.+)"
-+
-+            rest = match_v.groups()[0]
-+
-+            if match_layer := re.match(deepstack_merger_pattern, rest):
-+                layer_idx, layer_rest = match_layer.groups()
-+                layer_idx = str(int(layer_idx) + offsets)
-+                hf_name_param = [(f"model.visual.deepstack_merger_list.{layer_idx}.{layer_rest}", param)]
-+
-+            # Decoder layers
-+            elif match_layer := re.match(decoder_layer_pattern, rest):
-+
-+                layer_idx, layer_rest = match_layer.groups()
-+                layer_idx = str(int(layer_idx) + offsets)
-+                hf_name_param = [(f"model.visual.blocks.{layer_idx}.{layer_rest}", param)]
-+            else:
-+                hf_name_param = [(f"model.visual.{rest}", param)]
-+
-+    if hf_name_param == None:
-+        raise ValueError(f"Unknown parameter name: {name}")
-+    else:
-+        return hf_name_param
-diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py
-index 8174c7ac..1f1c0e09 100644
---- a/slime/backends/megatron_utils/model_provider.py
-+++ b/slime/backends/megatron_utils/model_provider.py
-@@ -17,7 +17,7 @@ from megatron.core.transformer.transformer_config import TransformerConfig
- from megatron.training.arguments import core_transformer_config_from_args
- 
- from slime.utils.misc import load_function
--
-+import slime_plugins.patch.mbridge_patch
- 
- # Adapt from https://github.com/volcengine/verl/blob/c3b20575d2bc815fcccd84bddb4c0401fc4b632b/verl/models/llama/megatron/layers/parallel_linear.py#L82
- class LinearForLastLayer(torch.nn.Linear):
-@@ -33,7 +33,7 @@ class LinearForLastLayer(torch.nn.Linear):
-         self.sequence_parallel = config.sequence_parallel
-         if self.sequence_parallel:
-             self.weight.sequence_parallel = True
--
-+            self.bias.sequence_parallel = True
-         self.weight.data.normal_(mean=0.0, std=0.02)
-         if bias:
-             self.bias.data.zero_()
-@@ -51,10 +51,59 @@ class LinearForLastLayer(torch.nn.Linear):
-         return logits, None
- 
- 
-+def get_qwen3vl_provide_wrapper(provider, role):
-+
-+    def provide_wrapper(pre_process=None, post_process=None, vp_stage=None):
-+        """
-+        Provide a Qwen3VL MoE model instance with vision and language components.
-+        """
-+        from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel
-+        language_transformer_config = provider
-+
-+        # Create vision transformer config - placeholder for future use
-+        hf_config = provider.vision_config
-+
-+        language_transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(
-+            num_experts=provider.num_moe_experts,
-+            moe_grouped_gemm=True,
-+            qk_layernorm=provider.qk_layernorm,
-+            fp8=False,
-+            normalization="RMSNorm",
-+        )
-+
-+        # reuse Qwen3VLModel for MoE model but replace the language model with MoE language model
-+        model = Qwen3VLModel(
-+            language_transformer_config=language_transformer_config,
-+            language_transformer_layer_spec=language_transformer_layer_spec,
-+            vision_transformer_config=hf_config,
-+            pre_process=pre_process,
-+            post_process=post_process,
-+        )
-+
-+        if role == "critic" and post_process:
-+            model.language_model.output_layer = LinearForLastLayer(input_size=provider.hidden_size, output_size=1, config=provider).to(
-+                device=model.language_model.output_layer.weight.device,
-+                dtype=model.language_model.output_layer.weight.dtype,
-+            )
-+
-+        # Apply freeze options if any are enabled for fine-tuning
-+        if provider.freeze_language_model or provider.freeze_vision_model or provider.freeze_vision_projection:
-+            model.freeze(
-+                freeze_language_model=provider.freeze_language_model,
-+                freeze_vision_model=provider.freeze_vision_model,
-+                freeze_vision_projection=provider.freeze_vision_projection,
-+            )
-+
-+        return model
-+    return provide_wrapper
-+
-+
- def get_model_provider_func(
-     args: argparse.Namespace,
-     role: Literal["actor", "critic"] = "actor",
- ):
-+    from megatron.bridge.models.conversion.param_mapping import AutoMapping
-+    AutoMapping.register_module_type('LinearForLastLayer', 'replicated')  # 或 'column' / 'replicated'
-     # Support custom model provider path (similar to --custom-rm-path for reward models)
-     if getattr(args, "custom_model_provider_path", None):
- 
-@@ -88,6 +137,27 @@ def get_model_provider_func(
-         provider.expert_model_parallel_size = args.expert_model_parallel_size
-         provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size
-         provider.sequence_parallel = args.sequence_parallel
-+        provider.gradient_accumulation_fusion = args.gradient_accumulation_fusion
-+        provider.recompute_granularity = args.recompute_granularity
-+        provider.recompute_method = args.recompute_method
-+        provider.recompute_num_layers = args.recompute_num_layers
-+        for key, value in vars(args).items():
-+            if hasattr(provider, key):
-+                continue
-+            setattr(provider, key, value)
-+
-+        is_qwen3vl = (
-+            hasattr(bridge.hf_pretrained, 'config') 
-+            and hasattr(bridge.hf_pretrained.config, 'model_type') 
-+            and 'qwen3_vl' in bridge.hf_pretrained.config.model_type.lower()
-+        )
-+
-+        if role == 'critic' and is_qwen3vl:
-+            from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
-+            from slime_plugins.patch.critic_patch import load_weights_hf_to_megatron_wrapper
-+            MegatronModelBridge.load_weights_hf_to_megatron = load_weights_hf_to_megatron_wrapper
-+            provider.provide = get_qwen3vl_provide_wrapper(provider, role)
-+            
-         provider.finalize()
-         return provider.provide
- 
-diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py
-index a2e4e129..ce40b446 100644
---- a/slime/backends/megatron_utils/update_weight/common.py
-+++ b/slime/backends/megatron_utils/update_weight/common.py
-@@ -10,7 +10,7 @@ from megatron.core.transformer.transformer_layer import get_transformer_layer_of
- 
- from slime.backends.megatron_utils.misc_utils import strip_param_name_prefix
- from slime.utils.types import ParamInfo
--
-+from slime.utils.common import is_npu
- 
- def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
-     """
-@@ -40,6 +40,8 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
-     if "linear_fc1.weight" in name:
-         param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
-         param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-+        if is_npu():
-+            partition_dim = 0
-     # this is bug in megatron's grouped moe.
-     if "linear_fc2.weight" in name:
-         if partition_dim == 0:
-@@ -102,6 +104,8 @@ def all_gather_params_async(
-             if "linear_fc1.weight" in info.name:
-                 param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
-                 param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-+                if is_npu():
-+                    partition_dim = 0
-             # this is bug in megatron's grouped moe.
-             if "linear_fc2.weight" in info.name:
-                 if partition_dim == 0:
-diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
-index a8e50e0e..b3c6ac24 100644
---- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
-+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
-@@ -12,6 +12,7 @@ from ray.actor import ActorHandle
- from tqdm import tqdm
- 
- from slime.utils.distributed_utils import get_gloo_group, init_process_group
-+from slime.utils.common import is_npu
- 
- from ..megatron_to_hf import convert_to_hf
- from .common import all_gather_param, named_params_and_buffers
-@@ -253,6 +254,7 @@ def connect_rollout_engines_from_distributed(
-         master_port = sock.getsockname()[1]
-     world_size = len(rollout_engines) * args.rollout_num_gpus_per_engine + 1
- 
-+    backend = "hccl" if is_npu() else "nccl"
-     refs = [
-         engine.init_weights_update_group.remote(
-             master_address,
-@@ -260,12 +262,12 @@ def connect_rollout_engines_from_distributed(
-             i * args.rollout_num_gpus_per_engine + 1,
-             world_size,
-             group_name,
--            backend="nccl",
-+            backend=backend,
-         )
-         for i, engine in enumerate(rollout_engines)
-     ]
-     model_update_groups = init_process_group(
--        backend="nccl",
-+        backend=backend,
-         init_method=f"tcp://{master_address}:{master_port}",
-         world_size=world_size,
-         rank=0,
-diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py
-index af0f5620..03ded2c8 100644
---- a/slime/backends/sglang_utils/sglang_engine.py
-+++ b/slime/backends/sglang_utils/sglang_engine.py
-@@ -15,6 +15,7 @@ from urllib3.exceptions import NewConnectionError
- 
- from slime.ray.ray_actor import RayActor
- from slime.utils.http_utils import get_host_info
-+from slime.utils.common import is_npu
- 
- logger = logging.getLogger(__name__)
- 
-@@ -33,7 +34,10 @@ def get_base_gpu_id(args, rank):
- 
- 
- def _to_local_gpu_id(physical_gpu_id: int) -> int:
--    cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
-+    if is_npu():
-+        cvd = os.environ.get("ASCEND_RT_VISIBLE_DEVICES")
-+    else:
-+        cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
-     if not cvd:
-         return physical_gpu_id  # no remapping
-     # CUDA_VISIBLE_DEVICES can be like "4,5,6,7"
-diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py
-index 4cac07a2..098bdf92 100644
---- a/slime/ray/actor_group.py
-+++ b/slime/ray/actor_group.py
-@@ -5,6 +5,7 @@ from ray.util.placement_group import PlacementGroup
- from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
- 
- from slime.ray.utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST
-+from slime.utils.common import is_npu
- 
- 
- class RayTrainGroup:
-@@ -87,19 +88,19 @@ class RayTrainGroup:
- 
-             actor_impl = FSDPTrainRayActor
- 
--        TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl)
--
-+        TrainRayActor = ray.remote(runtime_env={"env_vars": env_vars})(actor_impl)
-+        device_name = "NPU" if is_npu() else "GPU"
-         # Create worker actors
-         self._actor_handlers = []
-         master_addr, master_port = None, None
-         for rank in range(world_size):
-             actor = TrainRayActor.options(
-                 num_cpus=num_gpus_per_actor,
--                num_gpus=num_gpus_per_actor,
-                 scheduling_strategy=PlacementGroupSchedulingStrategy(
-                     placement_group=pg,
-                     placement_group_bundle_index=reordered_bundle_indices[rank],
-                 ),
-+                resources={device_name:num_gpus_per_actor}
-             ).remote(world_size, rank, master_addr, master_port)
-             if rank == 0:
-                 master_addr, master_port = ray.get(actor.get_master_addr_and_port.remote())
-diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py
-index eb232b16..963b4071 100644
---- a/slime/ray/placement_group.py
-+++ b/slime/ray/placement_group.py
-@@ -4,6 +4,7 @@ import socket
- import ray
- from ray.util.placement_group import placement_group
- from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
-+from slime.utils.common import is_npu
- 
- from .actor_group import RayTrainGroup
- from .rollout import RolloutManager
-@@ -11,10 +12,13 @@ from .rollout import RolloutManager
- logger = logging.getLogger(__name__)
- 
- 
--@ray.remote(num_gpus=1)
-+@ray.remote
- class InfoActor:
-     def get_ip_and_gpu_id(self):
--        return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
-+        if is_npu():
-+            return ray.util.get_node_ip_address(), ray.get_runtime_context().get_accelerator_ids()["NPU"][0]
-+        else:
-+            return ray.util.get_node_ip_address(), ray.get_gpu_ids()[0]
- 
- 
- def sort_key(x):
-@@ -35,12 +39,13 @@ def sort_key(x):
-             # representation that allows for sorting.
-             node_ip_parts = [ord(c) for c in node_identifier]
- 
--    return (node_ip_parts, gpu_id)
-+    return (node_ip_parts, int(gpu_id))
- 
- 
- def _create_placement_group(num_gpus):
-     """Create a placement group with the specified number of GPUs."""
--    bundles = [{"GPU": 1, "CPU": 1} for _ in range(num_gpus)]
-+    device_name = "NPU" if is_npu() else "GPU"
-+    bundles = [{device_name: 1, "CPU": 1} for _ in range(num_gpus)]
-     pg = placement_group(bundles, strategy="PACK")
-     num_bundles = len(bundles)
- 
-@@ -53,7 +58,8 @@ def _create_placement_group(num_gpus):
-                 scheduling_strategy=PlacementGroupSchedulingStrategy(
-                     placement_group=pg,
-                     placement_group_bundle_index=i,
--                )
-+                ),
-+                resources={device_name:1}
-             ).remote()
-         )
-     gpu_ids = ray.get([actor.get_ip_and_gpu_id.remote() for actor in info_actors])
-@@ -167,9 +173,11 @@ def create_training_models(args, pgs, rollout_manager):
- 
- 
- def create_rollout_manager(args, pg):
-+    device_name = "NPU" if is_npu() else "GPU"
-     rollout_manager = RolloutManager.options(
-         num_cpus=1,
-         num_gpus=0,
-+        resources={device_name:0}
-     ).remote(args, pg)
- 
-     # calculate num_rollout from num_epoch
-diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py
-index 75cb053c..c54a3854 100644
---- a/slime/ray/rollout.py
-+++ b/slime/ray/rollout.py
-@@ -28,6 +28,7 @@ from slime.utils.metric_utils import (
- from slime.utils.misc import Box, group_by, load_function
- from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
- from slime.utils.types import Sample
-+from slime.utils.common import is_npu
- 
- from ..utils.metric_utils import has_repetition
- from .utils import NOSET_VISIBLE_DEVICES_ENV_VARS_LIST, Lock
-@@ -76,7 +77,8 @@ class RolloutManager:
-             self.all_rollout_engines = [None] * num_engines
-         self.num_new_engines = init_rollout_engines(args, pg, self.all_rollout_engines)
-         self.nodes_per_engine = max(1, args.rollout_num_gpus_per_engine // args.num_gpus_per_node)
--        self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()
-+        device_name = "NPU" if is_npu() else "GPU"
-+        self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0, resources={device_name:0}).remote()
-         self.rollout_id = -1
- 
-         self._metric_checker = MetricChecker.maybe_create(args)
-@@ -467,6 +469,7 @@ def init_rollout_engines(args, pg, all_rollout_engines):
-     RolloutRayActor = ray.remote(SGLangEngine)
- 
-     rollout_engines = []
-+    device_name = "NPU" if is_npu() else "GPU"
-     for i in range(num_engines):
-         if all_rollout_engines[i] is not None:
-             continue
-@@ -503,11 +506,11 @@ def init_rollout_engines(args, pg, all_rollout_engines):
- 
-         rollout_engine = RolloutRayActor.options(
-             num_cpus=num_cpus,
--            num_gpus=num_gpus,
-             scheduling_strategy=scheduling_strategy,
-             runtime_env={
-                 "env_vars": env_vars,
-             },
-+            resources={device_name:num_gpus}
-         ).remote(args, rank=i, worker_type=worker_type, base_gpu_id=base_gpu_id)
- 
-         rollout_engines.append((i, rollout_engine))
-diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py
-index 2e900ca5..d0a25583 100644
---- a/slime/ray/train_actor.py
-+++ b/slime/ray/train_actor.py
-@@ -13,16 +13,23 @@ from slime.ray.ray_actor import RayActor
- from slime.utils.distributed_utils import init_gloo_group
- from slime.utils.logging_utils import configure_logger
- from slime.utils.memory_utils import clear_memory, print_memory
-+from slime.utils.common import is_npu
- 
- logger = logging.getLogger(__name__)
- 
- 
- def get_local_gpu_id():
--    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", None)
-+    if is_npu():
-+        env_var = "ASCEND_RT_VISIBLE_DEVICES"
-+        device_ids = ray.get_runtime_context().get_accelerator_ids()["NPU"]
-+    else:
-+        env_var = "CUDA_VISIBLE_DEVICES"
-+        device_ids = ray.get_gpu_ids()
-+    cvd = os.environ.get(env_var, None)
-     if cvd is None:
--        return ray.get_gpu_ids()[0]
-+        return device_ids[0]
-     else:
--        return cvd.split(",").index(str(ray.get_gpu_ids()[0]))
-+        return cvd.split(",").index(str(device_ids[0]))
- 
- 
- class TrainRayActor(RayActor):
-diff --git a/slime/utils/common.py b/slime/utils/common.py
-new file mode 100644
-index 00000000..60ee5b5c
---- /dev/null
-+++ b/slime/utils/common.py
-@@ -0,0 +1,12 @@
-+import torch
-+
-+def is_npu() -> bool:
-+    if not hasattr(torch, "npu"):
-+        return False
-+
-+    if not torch.npu.is_available():
-+        raise RuntimeError(
-+            "torch_npu detected, but NPU device is not available or visible."
-+        )
-+
-+    return True
-diff --git a/slime/utils/external_utils/command_utils.py b/slime/utils/external_utils/command_utils.py
-index 9f51ecdf..d4b47eca 100644
---- a/slime/utils/external_utils/command_utils.py
-+++ b/slime/utils/external_utils/command_utils.py
-@@ -184,6 +184,108 @@ def execute_train(
-             f"{train_args}"
-         )
- 
-+def execute_train_npu(
-+    train_args: str,
-+    megatron_model_type: str | None,
-+    train_script: str = "train.py",
-+    before_ray_job_submit=None,
-+    extra_env_vars=None,
-+    config: ExecuteTrainConfig | None = None,
-+):
-+    if extra_env_vars is None:
-+        extra_env_vars = {}
-+    if config is None:
-+        config = ExecuteTrainConfig()
-+    external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY")
-+    master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
-+
-+    train_backend_fsdp = "--train-backend fsdp" in train_args
-+    assert train_backend_fsdp == (megatron_model_type is None)
-+
-+    exec_command(
-+        "pkill -9 sglang; "
-+        "sleep 3; "
-+        f"{'' if external_ray else 'ray stop --force; '}"
-+        f"{'' if external_ray else 'pkill -9 ray; '}"
-+        # cannot be run in CI, o/w kill the parent script
-+        # TODO: do we really need this kill? (or can we instead kill slime)
-+        # "pkill -9 python; "
-+        "pkill -9 slime; "
-+        "sleep 3; "
-+        f"{'' if external_ray else 'pkill -9 ray; '}"
-+        # "pkill -9 python; "
-+        "pkill -9 slime; "
-+        "pkill -9 redis; "
-+        "true; "
-+    )
-+
-+    if not external_ray:
-+        exec_command(
-+            # will prevent ray from buffering stdout/stderr
-+            f"export PYTHONBUFFERED=16 && "
-+            f"ray start --head --node-ip-address {master_addr} --disable-usage-stats"
-+        )
-+
-+    if (f := before_ray_job_submit) is not None:
-+        f()
-+
-+    runtime_env_json = json.dumps(
-+        {
-+            "env_vars": {
-+                # "PYTHONPATH": "/root/Megatron-LM/",
-+                "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-+                "RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES": "1",
-+                "ASCEND_TOOLKIT_HOME": "/path/to/ascend/Ascend/ascend-toolkit/latest/",
-+                "ASCEND_OPP_PATH": "/path/to/ascend/Ascend/ascend-toolkit/latest/opp/",
-+                "ASCEND_AICPU_PATH": "/path/to/ascend/Ascend/ascend-toolkit/latest/",
-+                "ASCEND_HOME_PATH": "/path/to/ascend/Ascend/ascend-toolkit/latest/",
-+                "set_env_path": "/path/to/ascend/Ascend/nnal/atb/set_env.sh",
-+                "HYDRA_FULL_ERROR": "1",
-+                "HCCL_HOST_SOCKET_PORT_RANGE": "60000-60050",
-+                "HCCL_NPU_SOCKET_PORT_RANGE": "61000-61050",
-+                # If setting this in FSDP, the computation communication overlapping may have issues
-+                **(
-+                    {}
-+                    if train_backend_fsdp
-+                    else {
-+                        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-+                    }
-+                ),
-+                "NCCL_NVLS_ENABLE": str(int(check_has_nvlink())),
-+                "no_proxy": f"127.0.0.1,{master_addr}",
-+                # This is needed by megatron / torch distributed in multi-node setup
-+                "MASTER_ADDR": master_addr,
-+                **(
-+                    {
-+                        "CUDA_ENABLE_COREDUMP_ON_EXCEPTION": "1",
-+                        "CUDA_COREDUMP_SHOW_PROGRESS": "1",
-+                        "CUDA_COREDUMP_GENERATION_FLAGS": "skip_nonrelocated_elf_images,skip_global_memory,skip_shared_memory,skip_local_memory,skip_constbank_memory",
-+                        "CUDA_COREDUMP_FILE": "/root/shared_data/cuda_coredump_%h.%p.%t",
-+                    }
-+                    if config.cuda_core_dump
-+                    else {}
-+                ),
-+                **extra_env_vars,
-+                **_parse_extra_env_vars(config.extra_env_vars),
-+            }
-+        }
-+    )
-+
-+    if get_bool_env_var("SLIME_SCRIPT_ENABLE_RAY_SUBMIT", "1"):
-+        cmd_megatron_model_source = (
-+            f'source "{repo_base_dir}/scripts/models/{megatron_model_type}.sh" && '
-+            if megatron_model_type is not None
-+            else ""
-+        )
-+        exec_command(
-+            f"export no_proxy=127.0.0.1 && export PYTHONBUFFERED=16 && "
-+            f"{cmd_megatron_model_source}"
-+            f'ray job submit --address="http://127.0.0.1:8265" '
-+            f"--runtime-env-json='{runtime_env_json}' "
-+            f"-- python3 {train_script} "
-+            f"{'${MODEL_ARGS[@]}' if megatron_model_type is not None else ''} "
-+            f"{train_args}"
-+        )
- 
- def _parse_extra_env_vars(text: str):
-     try:
-diff --git a/slime/utils/memory_utils.py b/slime/utils/memory_utils.py
-index c12f3cd0..89078266 100644
---- a/slime/utils/memory_utils.py
-+++ b/slime/utils/memory_utils.py
-@@ -3,6 +3,7 @@ import logging
- 
- import torch
- import torch.distributed as dist
-+from slime.utils.common import is_npu
- 
- logger = logging.getLogger(__name__)
- 
-@@ -12,12 +13,19 @@ def clear_memory(clear_host_memory: bool = False):
-     gc.collect()
-     torch.cuda.empty_cache()
-     if clear_host_memory:
--        torch._C._host_emptyCache()
-+        if is_npu():
-+            torch.npu.empty_cache()
-+        else:
-+            torch._C._host_emptyCache()
- 
- 
- def available_memory():
--    device = torch.cuda.current_device()
--    free, total = torch.cuda.mem_get_info(device)
-+    if is_npu():
-+        device = torch.npu.current_device()
-+        free, total = torch.npu.mem_get_info(device)
-+    else:
-+        device = torch.cuda.current_device()
-+        free, total = torch.cuda.mem_get_info(device)
-     return {
-         "gpu": str(device),
-         "total_GB": _byte_to_gb(total),
-diff --git a/slime_plugins/patch/critic_patch.py b/slime_plugins/patch/critic_patch.py
-new file mode 100644
-index 00000000..1f4e1bc6
---- /dev/null
-+++ b/slime_plugins/patch/critic_patch.py
-@@ -0,0 +1,94 @@
-+import torch
-+from typing import (
-+    List,
-+    Mapping,
-+    TypeVar,
-+    Union,
-+)
-+from megatron.core.transformer.module import MegatronModule
-+HFPreTrained = TypeVar("HFPreTrained")
-+MegatronModel = TypeVar("MegatronModel", bound=MegatronModule)
-+
-+
-+def load_weights_hf_to_megatron_wrapper(
-+        self, hf_pretrained: HFPreTrained, megatron_model: Union[MegatronModel, List[MegatronModel]]
-+    ) -> List[MegatronModel]:
-+        """Load HuggingFace weights into Megatron models.
-+
-+        This method orchestrates the complete weight loading process from HuggingFace
-+        format to Megatron's distributed format. It builds a conversion task and
-+        executes it with proper progress tracking and error handling.
-+
-+        The actual weight transformations and distribution are delegated to the
-+        appropriate MegatronParamMapping instances based on the state mappings.
-+
-+        Args:
-+            hf_pretrained (HFPreTrained): HuggingFace model or state source containing the
-+                weights to load.
-+            megatron_model (Union[MegatronModel, List[MegatronModel]]): Megatron model instance
-+                or list of model instances (one per virtual pipeline stage).
-+
-+        Returns:
-+            List[MegatronModel]: The input megatron_model as a list with loaded weights.
-+
-+        Process:
-+        1. Build a task mapping each Megatron parameter to its source
-+        2. For each parameter in the task:
-+            - Fetch source weights from HuggingFace state
-+            - Apply format transformation via the param mapping
-+            - Distribute to appropriate TP/PP ranks
-+            - Copy into the Megatron parameter
-+
-+        Example:
-+            .. code-block:: python
-+
-+                hf_model = PreTrainedCausalLM.from_pretrained("gpt2")
-+                megatron_model = create_megatron_model()  # Single model or list
-+                bridge.load_weights_hf_to_megatron(hf_model, megatron_model)
-+
-+        Note:
-+            Progress is shown only on rank 0 to avoid cluttered output in
-+            distributed environments.
-+
-+        Raises:
-+            ValueError: If hf_pretrained doesn't have state attribute or if weight shapes don't match.
-+            AttributeError: If required HF weights are missing.
-+        """
-+        if not isinstance(megatron_model, list):
-+            megatron_model = [megatron_model]
-+
-+        hf_to_megatron_tasks = self.build_conversion_tasks(hf_pretrained, megatron_model)
-+        hf_state_dict: Mapping[str, torch.Tensor] = hf_pretrained.state if hasattr(hf_pretrained, "state") else {}
-+
-+        description = f"Loading from {hf_pretrained.model_name_or_path}"
-+        for task in self._with_progress_tracking(hf_to_megatron_tasks, description):
-+            # None means megatron module not on current rank, skip if this task is not going to happen
-+            if task.megatron_module is None:
-+                continue
-+            # 1) Fetch source tensor(s) from HF state dict
-+            hf_weights = self.maybe_modify_loaded_hf_weight(task.mapping.hf_param, hf_state_dict)
-+
-+            # 2) Delegate conversion & distribution to the bridge
-+            converted_weights = task.mapping.hf_to_megatron(hf_weights, task.megatron_module)
-+
-+            # 3) Copy into Megatron param if this rank received a shard
-+            if converted_weights is not None:
-+                # Assert that param_weight is not None for HF->Megatron tasks
-+                assert task.param_weight is not None, "param_weight is required for HF->Megatron conversion"
-+
-+                if converted_weights.shape != task.param_weight.shape and task.param_name == 'language_model.output_layer.weight':
-+                    continue
-+
-+                # Check shape compatibility before copying
-+                if converted_weights.shape != task.param_weight.shape:
-+                    raise ValueError(
-+                        f"Shape mismatch for megatron param {task.mapping.megatron_param}:\n"
-+                        f"  Expected shape: {task.param_weight.shape}\n"
-+                        f"  Got shape: {converted_weights.shape}\n"
-+                        f"  Bridge type: {type(task.mapping).__name__}\n"
-+                        f"  HF mapping: {task.mapping.hf_param}"
-+                    )
-+                task.param_weight.data.copy_(converted_weights)
-+
-+        self._broadcast_shared_embeddings(megatron_model)
-+        return megatron_model
-\ No newline at end of file
-diff --git a/slime_plugins/patch/mbridge_patch.py b/slime_plugins/patch/mbridge_patch.py
-new file mode 100644
-index 00000000..606fff9b
---- /dev/null
-+++ b/slime_plugins/patch/mbridge_patch.py
-@@ -0,0 +1,25 @@
-+try:
-+    from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
-+    _original_build_conversion_tasks = MegatronModelBridge.build_conversion_tasks
-+
-+    def _patched_build_conversion_tasks(self, hf_pretrained, megatron_model):
-+        """
-+        Invoke the original build_conversion_tasks and filter out any actual None tasks.
-+
-+        The original implementation might return List[None | WeightConversionTask]. 
-+        We consolidate it here into List[WeightConversionTask] to avoid errors 
-+        when accessing None.task.xxx later.
-+        """
-+        tasks = _original_build_conversion_tasks(self, hf_pretrained, megatron_model)
-+
-+        if tasks is None:
-+            return []
-+
-+        filtered = [t for t in tasks if t is not None]
-+
-+        return filtered
-+
-+    MegatronModelBridge.build_conversion_tasks = _patched_build_conversion_tasks
-+
-+except ImportError:
-+    pass
-diff --git a/train.py b/train.py
-index 01883c47..18faa6d2 100644
---- a/train.py
-+++ b/train.py
-@@ -1,5 +1,7 @@
- import ray
--
-+from slime.utils.common import is_npu
-+if is_npu():
-+    import mindspeed.megatron_adaptor
- from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
- from slime.utils.arguments import parse_args
- from slime.utils.logging_utils import configure_logger, init_tracking
diff --git a/docker/patch/gb10/cuda_profiler_api.h b/docker/patch/gb10/cuda_profiler_api.h
deleted file mode 100644
index 276b80b656..0000000000
--- a/docker/patch/gb10/cuda_profiler_api.h
+++ /dev/null
@@ -1,26 +0,0 @@
-#ifndef __CUDA_PROFILER_API_H__
-#define __CUDA_PROFILER_API_H__
-
-/* CUDA 13 dropped the public cuda_profiler_api.h header, but the runtime still
- * exports cudaProfilerStart / cudaProfilerStop from libcudart.so. This shim
- * declares them against the public types so legacy callers (e.g., TE 2.10,
- * which only includes the header without actually calling the APIs in the
- * affected TUs) continue to compile.
- *
- * Installed by slime GB10 port. See docker/patch/gb10/ and NOTES_GB10.md.
- */
-
-#include <cuda_runtime_api.h>
-
-#if defined(__cplusplus)
-extern "C" {
-#endif
-
-extern __host__ cudaError_t CUDARTAPI cudaProfilerStart(void);
-extern __host__ cudaError_t CUDARTAPI cudaProfilerStop(void);
-
-#if defined(__cplusplus)
-}
-#endif
-
-#endif /* __CUDA_PROFILER_API_H__ */
diff --git a/docker/patch/gb10/patch_sgl_kernel.py b/docker/patch/gb10/patch_sgl_kernel.py
deleted file mode 100644
index a6dc3bbebf..0000000000
--- a/docker/patch/gb10/patch_sgl_kernel.py
+++ /dev/null
@@ -1,52 +0,0 @@
-"""Patch sgl-kernel CMakeLists.txt to add a SGL_KERNEL_GB10_ONLY build option.
-
-On GB10 (DGX Spark, sm_121a aarch64 with unified 128 GB memory), the stock
-sgl-kernel CMake emits gencodes for sm_90a + sm_100a + sm_103a + sm_110a +
-sm_120a + sm_121a. Cutlass FP8 gemm template instantiation per extra arch uses
-10-15 GB RAM, and 12 parallel nvcc jobs OOM-kill cicc.
-
-This patch wraps the non-GB10 gencode blocks in `if (NOT SGL_KERNEL_GB10_ONLY)`
-and adds a GB10-only branch that emits just sm_120a + sm_121a. Default OFF
-preserves upstream behavior.
-"""
-
-import sys
-from pathlib import Path
-
-p = Path(sys.argv[1] if len(sys.argv) > 1 else "/root/src/sglang/sgl-kernel/CMakeLists.txt")
-s = p.read_text()
-
-marker = 'option(SGL_KERNEL_ENABLE_SM100A           "Enable SM100A"           OFF)'
-new_opt = (
-    marker
-    + "\n"
-    + 'option(SGL_KERNEL_GB10_ONLY              "Build only for GB10 (sm_121a + sm_120a). Skips sm_90a/sm_100a/sm_103a/sm_110a and FA3." OFF)'
-)
-assert marker in s, "marker 1 (SM100A option line) not found"
-s = s.replace(marker, new_opt, 1)
-
-old = 'if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)'
-assert old in s, "marker 2 (CUDA 12.8 block) not found"
-s = s.replace(old, "if (NOT SGL_KERNEL_GB10_ONLY)\n\n" + old, 1)
-
-old2 = 'if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4)'
-assert old2 in s, "marker 3 (FP4 block) not found"
-insert_before = """endif()  # NOT SGL_KERNEL_GB10_ONLY
-
-if (SGL_KERNEL_GB10_ONLY)
-    # GB10 (DGX Spark) is sm_121a (consumer Blackwell, aarch64). It cannot run
-    # Hopper (sm_90a) nor datacenter Blackwell (sm_100a) binaries. Compiling
-    # them wastes ~6x build time and ~10-15 GB RAM per TU (cutlass templates),
-    # which OOM-kills cicc on 128 GB Spark hosts.
-    list(APPEND SGL_KERNEL_CUDA_FLAGS
-        "-gencode=arch=compute_120a,code=sm_120a"
-        "-gencode=arch=compute_121a,code=sm_121a"
-        "--compress-mode=size"
-    )
-endif()
-
-"""
-s = s.replace(old2, insert_before + old2, 1)
-
-p.write_text(s)
-print(f"edits applied to {p}")
diff --git a/docker/patch/gb10/sgl-kernel-arch.patch b/docker/patch/gb10/sgl-kernel-arch.patch
deleted file mode 100644
index a507e04c09..0000000000
--- a/docker/patch/gb10/sgl-kernel-arch.patch
+++ /dev/null
@@ -1,40 +0,0 @@
---- CMakeLists.txt.orig	2026-04-15 07:39:19.526964507 +0000
-+++ CMakeLists.txt	2026-04-15 07:39:19.548828896 +0000
-@@ -188,6 +188,7 @@
- option(SGL_KERNEL_ENABLE_FA3              "Enable FA3"              OFF)
- option(SGL_KERNEL_ENABLE_SM90A            "Enable SM90A"            OFF)
- option(SGL_KERNEL_ENABLE_SM100A           "Enable SM100A"           OFF)
-+option(SGL_KERNEL_GB10_ONLY              "Build only for GB10 (sm_121a + sm_120a). Skips sm_90a/sm_100a/sm_103a/sm_110a and FA3." OFF)
- 
- if (SGL_KERNEL_ENABLE_BF16)
-     list(APPEND SGL_KERNEL_CUDA_FLAGS
-@@ -216,6 +217,8 @@
- 
- endif()
- 
-+if (NOT SGL_KERNEL_GB10_ONLY)
-+
- if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_SM100A)
-     list(APPEND SGL_KERNEL_CUDA_FLAGS
-         "-gencode=arch=compute_100a,code=sm_100a"
-@@ -249,6 +252,20 @@
-     )
- endif()
- 
-+endif()  # NOT SGL_KERNEL_GB10_ONLY
-+
-+if (SGL_KERNEL_GB10_ONLY)
-+    # GB10 (DGX Spark) is sm_121a (consumer Blackwell, aarch64). It cannot run
-+    # Hopper (sm_90a) nor datacenter Blackwell (sm_100a) binaries. Compiling
-+    # them wastes ~6x build time and ~10-15 GB RAM per TU (cutlass templates),
-+    # which OOM-kills cicc on 128 GB Spark hosts.
-+    list(APPEND SGL_KERNEL_CUDA_FLAGS
-+        "-gencode=arch=compute_120a,code=sm_120a"
-+        "-gencode=arch=compute_121a,code=sm_121a"
-+        "--compress-mode=size"
-+    )
-+endif()
-+
- if ("${CUDA_VERSION}" VERSION_GREATER_EQUAL "12.8" OR SGL_KERNEL_ENABLE_FP4)
-     list(APPEND SGL_KERNEL_CUDA_FLAGS
-         "-DENABLE_NVFP4=1"
diff --git a/docker/patch/latest/megatron.patch b/docker/patch/latest/megatron.patch
index 0fb483a4ea..6d2a233949 100644
--- a/docker/patch/latest/megatron.patch
+++ b/docker/patch/latest/megatron.patch
@@ -12,10 +12,10 @@ index 41c21d93d..ef80f72d6 100644
              err_msg = f'Common file {load_path} does not exist'
              if MultiStorageClientFeature.is_enabled():
 diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py
-index a5b6c009b..22794d7e6 100644
+index 5a1ea308d..aa701237f 100644
 --- a/megatron/core/dist_checkpointing/strategies/torch.py
 +++ b/megatron/core/dist_checkpointing/strategies/torch.py
-@@ -503,10 +503,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
+@@ -597,10 +597,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
      def _validate_global_shapes(self, metadata, sharded_tensors):
          for sh_ten in sharded_tensors:
              if sh_ten.key not in metadata.state_dict_metadata:
@@ -30,9 +30,9 @@ index a5b6c009b..22794d7e6 100644
 +                print(f"{sh_ten.key} from model not in state dict, will skip")
 +                continue
              loaded_shape = metadata.state_dict_metadata[sh_ten.key].size
-             expected_shape = sh_ten.global_shape
+             expected_shape = self._expected_shape(sh_ten)
              if loaded_shape != expected_shape:
-@@ -530,7 +532,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
+@@ -630,7 +632,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
          tensor_metadata = self.metadata.state_dict_metadata
          metadata_with_sizes = [
              (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor)
@@ -41,7 +41,7 @@ index a5b6c009b..22794d7e6 100644
          ]
          try:
              # Temporarily set sizes to expected shapes
-@@ -802,6 +804,7 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
+@@ -959,6 +961,7 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
              planner=MCoreLoadPlanner(
                  shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
                  allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,
@@ -50,18 +50,18 @@ index a5b6c009b..22794d7e6 100644
          )
  
 diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
-index ef8527e9e..57fbe5bd7 100644
+index acb93ef78..d239db4ab 100644
 --- a/megatron/core/extensions/transformer_engine.py
 +++ b/megatron/core/extensions/transformer_engine.py
-@@ -639,6 +639,7 @@ class TELinear(te.pytorch.Linear):
-         self.te_quant_params: Optional[TEQuantizationParams] = None
+@@ -408,6 +408,7 @@ class TELinear(te.pytorch.Linear):
+         )
  
          for param in self.parameters():
 +            setattr(param, "parallel_mode", parallel_mode)
              if is_expert:
                  # Reduce the gradient on the expert_data_parallel group for expert linear layers
                  setattr(param, "allreduce", not self.expert_parallel)
-@@ -1455,6 +1456,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
+@@ -1161,6 +1162,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
  
  
  if HAVE_TE and is_te_min_version("1.9.0.dev0"):
@@ -115,15 +115,23 @@ index ef8527e9e..57fbe5bd7 100644
 +
 +    def fake_int4_quantization_ste(x, group_size):
 +        x_out = _FakeInt4QuantizationSTE.apply(x, group_size)
-+
++        
 +        if hasattr(x, 'main_grad'):
 +            x_out.main_grad = x.main_grad
-+
++            
 +        return x_out
  
      class TEGroupedLinear(te.pytorch.GroupedLinear):
          """
-@@ -1671,6 +1727,20 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"):
+@@ -1351,6 +1407,7 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"):
+             _is_first_microbatch = (
+                 None if self.disable_parameter_transpose_cache else self.is_first_microbatch
+             )
++
+             out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
+             self.is_first_microbatch = False
+ 
+@@ -1361,6 +1418,20 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"):
                  return out
              return out, None
  
@@ -135,17 +143,17 @@ index ef8527e9e..57fbe5bd7 100644
 +                group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128"))
 +
 +                weight_tensors = [
-+                    fake_int4_quantization_ste(w, group_size)
++                    fake_int4_quantization_ste(w, group_size) 
 +                    for w in weight_tensors
 +                ]
-+
++                
 +            return weight_tensors
 +
          def _encode_extra_state(self, state):
              # TE 2.0 changed the format of extra_state to be a byte tensor
              if is_te_min_version("2.0.0"):
 diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py
-index 1fd5dcfae..75e1072d5 100644
+index 1fd5dcfae..c9aeef1f0 100644
 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py
 +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py
 @@ -385,6 +385,7 @@ def rotary_fwd_kv_kernel(
@@ -244,7 +252,7 @@ index 1fd5dcfae..75e1072d5 100644
 -    tl.store(dKV_ptr + dk_out_off, dk, mask=mask)
 -    tl.store(dKV_ptr + dv_out_off, dv, mask=mask)
 +    dKV_ptr = dKV + pid_m * stride_dkv_seq # + pid_head * BLOCK_H * stride_dkv_nheads
-+    ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H
++    ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H 
 +    kj_range = tl.arange(0, k_dim_ceil)[None, :]
 +    mask_k = (ki_range < head_num) & (kj_range < k_dim)
 +    mask_v = ki_range < head_num
@@ -256,7 +264,7 @@ index 1fd5dcfae..75e1072d5 100644
 +
 +    dk = tl.load(dK_ptr + dk_in_off, mask=mask_k)
 +    tl.store(dKV_ptr + dk_out_off, dk, mask=mask_k)
-+
++    
 +    if v_dim > 0:
 +        dv_out_off = ki_range * stride_dkv_nheads + k_dim + tl.arange(0, v_dim)[None, :]
 +        dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :]
@@ -306,121 +314,83 @@ index 1fd5dcfae..75e1072d5 100644
              ctx.v_dim,
              nheads,
              batch_size,
-diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py
-index 5dc2d5030..2b241c86d 100644
---- a/megatron/core/inference/contexts/dynamic_context.py
-+++ b/megatron/core/inference/contexts/dynamic_context.py
-@@ -61,8 +61,8 @@ except ImportError:
- try:
-     from torch_memory_saver import torch_memory_saver
- 
--    torch_memory_saver.hook_mode = "torch"
--    HAVE_TORCH_MEMORY_SAVER = True
-+    # torch_memory_saver.hook_mode = "torch"
-+    HAVE_TORCH_MEMORY_SAVER = False
- except ImportError:
-     HAVE_TORCH_MEMORY_SAVER = False
- 
-diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-index 05a7e8f60..881cfbcaa 100644
---- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-+++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-@@ -308,6 +308,7 @@ class MultimodalRotaryEmbedding(nn.Module):
-         self,
-         position_ids: torch.Tensor,
-         mrope_section: List[int],
-+        packed_seq: bool = False,
-         cp_group: Optional[torch.distributed.ProcessGroup] = None,
-     ) -> Tensor:
-         """Forward pass of multimodal RoPE embedding.
-@@ -352,7 +353,9 @@ class MultimodalRotaryEmbedding(nn.Module):
-         emb = emb[..., None, :].transpose(0, 1).contiguous()
-         if cp_group is None:
-             cp_group = self.cp_group
--        if cp_group is not None and cp_group.size() > 1:
-+        # For THD (packed sequence) format, skip CP slicing here — it is handled
-+        # per-sequence inside _apply_rotary_pos_emb_thd instead (same as RotaryEmbedding).
-+        if cp_group is not None and cp_group.size() > 1 and not packed_seq:
-             # slice rotary_pos_emb along sequence dimension and select the parition of the current
-             # CP rank
-             emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
+diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
+index 13d74aa52..060898a7a 100644
+--- a/megatron/core/models/common/language_module/language_module.py
++++ b/megatron/core/models/common/language_module/language_module.py
+@@ -184,7 +184,15 @@ class LanguageModule(MegatronModule):
+             assert (
+                 column_parallel_linear is not None
+             ), "column_parallel_linear cannot be None when not using fused linear cross entropy."
+-            logits, _ = column_parallel_linear(hidden, **col_linear_kwargs)
++            # output
++            output_layer_params = {k: v.detach() for k, v in column_parallel_linear.named_parameters()}
++            output_layer_buffers = dict(column_parallel_linear.named_buffers())
++            logits, _ = torch.func.functional_call(
++                column_parallel_linear,
++                {**output_layer_params, **output_layer_buffers},
++                (hidden,),
++                col_linear_kwargs,
++            )
+ 
+             return self.compute_language_model_loss(labels, logits)
+ 
 diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
-index 5bb479ad3..a9d3583e5 100755
+index e21127b87..712793853 100755
 --- a/megatron/core/models/gpt/gpt_layer_specs.py
 +++ b/megatron/core/models/gpt/gpt_layer_specs.py
-@@ -182,6 +182,8 @@ def get_gpt_layer_with_transformer_engine_spec(
+@@ -188,6 +188,8 @@ def get_gpt_layer_with_transformer_engine_spec(
+     use_kitchen: bool = False,
+     use_te_activation_func: bool = False,
      fallback_to_eager_attn: bool = False,
-     use_kitchen_attention: bool = False,
-     kitchen_attention_backend: str = "sdpa",
 +    post_self_attn_layernorm: bool = False,
 +    post_mlp_layernorm: bool = False,
  ) -> ModuleSpec:
      """Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
  
-@@ -263,9 +265,11 @@ def get_gpt_layer_with_transformer_engine_spec(
-                     ),
-                 ),
-                 self_attn_bda=get_bias_dropout_add,
-+                post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
-                 pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp,
-                 mlp=mlp,
-                 mlp_bda=get_bias_dropout_add,
-+                post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
-             ),
-         )
-     else:
-@@ -289,9 +293,11 @@ def get_gpt_layer_with_transformer_engine_spec(
-                     ),
-                 ),
-                 self_attn_bda=get_bias_dropout_add,
-+                post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
-                 pre_mlp_layernorm=backend.layer_norm() if num_experts else IdentityOp,
-                 mlp=mlp,
-                 mlp_bda=get_bias_dropout_add,
-+                post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
-                 sharded_state_dict_keys_map={
-                     "mlp.0.weight": "mlp.linear_fc1.layer_norm_weight",
-                     "mlp.0.bias": "mlp.linear_fc1.layer_norm_bias",
-@@ -537,6 +543,8 @@ def get_gpt_decoder_layer_specs(
-             qk_l2_norm=qk_l2_norm,
-             use_kitchen=config.use_kitchen,
-             use_te_activation_func=config.use_te_activation_func,
-+            post_self_attn_layernorm=config.post_self_attn_layernorm,
-+            post_mlp_layernorm=config.post_mlp_layernorm,
-         )
-         moe_layer_spec = get_gpt_layer_with_transformer_engine_spec(
-             num_experts=config.num_moe_experts,
-@@ -547,6 +555,8 @@ def get_gpt_decoder_layer_specs(
-             qk_l2_norm=qk_l2_norm,
-             use_kitchen=config.use_kitchen,
-             use_te_activation_func=config.use_te_activation_func,
-+            post_self_attn_layernorm=config.post_self_attn_layernorm,
-+            post_mlp_layernorm=config.post_mlp_layernorm,
-         )
-     else:
-         dense_layer_spec = get_gpt_layer_local_spec(
+@@ -260,6 +262,8 @@ def get_gpt_layer_with_transformer_engine_spec(
+         mlp=mlp,
+         sharded_state_dict_keys_map=sharded_state_dict_keys_map,
+         normalization=normalization,
++        post_self_attn_layernorm=post_self_attn_layernorm,
++        post_mlp_layernorm=post_mlp_layernorm,
+     )
+ 
+ 
+@@ -349,6 +353,8 @@ def get_transformer_layer_spec_for_backend(
+     mlp: ModuleSpec,
+     sharded_state_dict_keys_map: Optional[dict] = None,
+     normalization: Optional[str] = None,
++    post_self_attn_layernorm: bool = False,
++    post_mlp_layernorm: bool = False,
+ ) -> ModuleSpec:
+     """Helper function to get module spec for TransformerLayer"""
+ 
+@@ -371,9 +377,11 @@ def get_transformer_layer_spec_for_backend(
+             input_layernorm=input_layernorm,
+             self_attention=attention,
+             self_attn_bda=get_bias_dropout_add,
++            post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
+             pre_mlp_layernorm=pre_mlp_layernorm,
+             mlp=mlp,
+             mlp_bda=get_bias_dropout_add,
++            post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
+             sharded_state_dict_keys_map=sharded_state_dict_keys_map,
+         ),
+     )
 diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
-index 5b31ddedf..ead60f2dd 100644
+index a1230568c..1fd52f65a 100644
 --- a/megatron/core/models/gpt/gpt_model.py
 +++ b/megatron/core/models/gpt/gpt_model.py
-@@ -394,6 +394,8 @@ class GPTModel(LanguageModule):
-                 rotary_pos_emb = self.rotary_pos_emb(
-                     position_ids,
-                     self.mrope_section,
-+                    packed_seq=packed_seq_params is not None
-+                    and packed_seq_params.qkv_format == 'thd',
-                     cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None,
-                 )
-             else:
-@@ -488,6 +490,7 @@ class GPTModel(LanguageModule):
+@@ -446,6 +446,7 @@ class GPTModel(LanguageModule):
+         *,
          inference_params: Optional[BaseInferenceContext] = None,
          loss_mask: Optional[Tensor] = None,
-         padding_mask: Optional[Tensor] = None,
-+        mtp_kwargs: Optional[dict] = None,
++        mtp_kwargs: Optional[dict] = {},
      ) -> Tensor:
          """Forward function of the GPT Model This function passes the input tensors
          through the embedding layer, and then the decoder and finally into the post
-@@ -560,6 +563,7 @@ class GPTModel(LanguageModule):
+@@ -508,6 +509,7 @@ class GPTModel(LanguageModule):
              runtime_gather_output=runtime_gather_output,
              extra_block_kwargs=extra_block_kwargs,
              inference_context=inference_context,
@@ -428,80 +398,59 @@ index 5b31ddedf..ead60f2dd 100644
          )
  
      def _postprocess(
-@@ -581,6 +585,7 @@ class GPTModel(LanguageModule):
+@@ -529,6 +531,7 @@ class GPTModel(LanguageModule):
          runtime_gather_output=None,
          extra_block_kwargs=None,
          inference_context=None,
-+        mtp_kwargs=None,
++        mtp_kwargs={},
      ):
          """Postprocesses decoder hidden states to generate logits or compute loss.
  
-@@ -592,10 +597,12 @@ class GPTModel(LanguageModule):
-             assert runtime_gather_output, "Inference must always gather TP logits"
- 
-         # logits and loss
-+        mtp_kwargs = mtp_kwargs or {}
-+        mtp_labels = mtp_kwargs.get('mtp_labels')
+@@ -543,7 +546,8 @@ class GPTModel(LanguageModule):
          output_weight = None
          if self.share_embeddings_and_output_weights:
              output_weight = self.shared_embedding_or_output_weight()
 -        if mtp_in_postprocess:
-+        if mtp_in_postprocess and mtp_labels is not None:
++
++        if mtp_in_postprocess and mtp_kwargs.get('mtp_labels', None) is not None:
              hidden_states = self.mtp(
                  input_ids=input_ids,
                  position_ids=position_ids,
-@@ -614,13 +621,35 @@ class GPTModel(LanguageModule):
-         if not self.post_process:
+@@ -563,13 +567,18 @@ class GPTModel(LanguageModule):
              return hidden_states
  
--        if self.config.mtp_num_layers is not None:
+         # Skip when mtp_num_layers is None or 0
+-        if self.config.mtp_num_layers:
 -            mtp_labels = labels.clone()
-+        if self.config.mtp_num_layers and mtp_labels is not None:
-+            mtp_labels = mtp_labels.clone()
-+            mtp_labels, _ = roll_tensor(
-+                mtp_labels,
-+                shifts=-1,
-+                dims=-1,
-+                cp_group=self.cp_group,
-+                packed_seq_params=packed_seq_params,
-+            )
++        if self.config.mtp_num_layers and mtp_kwargs.get('mtp_labels', None) is not None:
++            mtp_labels = mtp_kwargs['mtp_labels'].clone()
++            mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
++
              hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
              hidden_states = hidden_states_list[0]
              if loss_mask is None:
                  # if loss_mask is not provided, use all ones as loss_mask
                  loss_mask = torch.ones_like(mtp_labels)
 +            else:
-+                loss_mask, _ = roll_tensor(
-+                    loss_mask,
-+                    shifts=-1,
-+                    dims=-1,
-+                    cp_group=self.cp_group,
-+                    packed_seq_params=packed_seq_params,
-+                )
-+
-+            mtp_output_weight = output_weight
-+            if mtp_output_weight is None and self.output_layer.weight is not None:
-+                mtp_output_weight = self.output_layer.weight
-+            if mtp_output_weight is not None:
-+                mtp_output_weight = mtp_output_weight.detach()
-+
++                # Otherwise, roll the loss_mask to keep up with the mtp_labels
++                loss_mask, _ = roll_tensor(loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
              for mtp_layer_number in range(self.config.mtp_num_layers):
                  # Calc loss for the current Multi-Token Prediction (MTP) layers.
                  mtp_labels, _ = roll_tensor(
-@@ -641,7 +670,7 @@ class GPTModel(LanguageModule):
-                 # Compute mtp loss without storing logits to save memory.
-                 output_layer_kwargs = dict(
-                     input_=hidden_states_list[mtp_layer_number + 1],
--                    weight=output_weight,
-+                    weight=mtp_output_weight,
-                     runtime_gather_output=runtime_gather_output,
+@@ -595,7 +604,7 @@ class GPTModel(LanguageModule):
+                     sequence_parallel_enabled=self.output_layer.sequence_parallel,
+                     column_parallel_linear=self.output_layer,
+                     col_linear_kwargs={
+-                        'weight': output_weight,
++                        'weight': output_weight.detach() if output_weight else None,
+                         'runtime_gather_output': runtime_gather_output,
+                     },
                  )
-                 if self.fuse_linear_cross_entropy:
 diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
-index a4364f5e9..c76f6daac 100644
+index 6e093f96f..eac21a3ea 100644
 --- a/megatron/core/optimizer/distrib_optimizer.py
 +++ b/megatron/core/optimizer/distrib_optimizer.py
-@@ -686,6 +686,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
+@@ -677,6 +677,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
                  # TE FusedAdam will not accumulate step for empty param groups, so we need to
                  # align the step across param groups.
                  param_group["step"] = int(step)
@@ -510,7 +459,7 @@ index a4364f5e9..c76f6daac 100644
  
          # Grad scaler state.
          if self.grad_scaler:
-@@ -1667,6 +1669,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
+@@ -1646,6 +1648,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
                              if key == 'padding':
                                  tensors[key] = LocalNonpersistentObject(tensors[key])
                                  continue
@@ -520,7 +469,7 @@ index a4364f5e9..c76f6daac 100644
                                  tensors[key].shape,
                                  gbuf_local_start,
 diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py
-index 7bb964078..2fe9a8cdc 100644
+index a273002b9..4f821cfd5 100644
 --- a/megatron/core/parallel_state.py
 +++ b/megatron/core/parallel_state.py
 @@ -11,6 +11,7 @@ from typing import Callable, List, Optional
@@ -529,7 +478,7 @@ index 7bb964078..2fe9a8cdc 100644
  import torch
 +import torch.distributed as dist
  
- from .utils import GlobalMemoryBuffer, GlobalSymmetricMemoryBuffer, is_torch_min_version
+ from .utils import GlobalMemoryBuffer, is_torch_min_version
  
 diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py
 index ac839c21f..f18309217 100644
@@ -563,12 +512,12 @@ index ac839c21f..f18309217 100644
          ops.append(recv_next_op)
      if len(ops) > 0:
 diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
-index 75825cd37..445b3fb84 100644
+index 28cff06f5..58dc4bb70 100644
 --- a/megatron/core/transformer/moe/moe_utils.py
 +++ b/megatron/core/transformer/moe/moe_utils.py
-@@ -711,6 +711,9 @@ def topk_routing_with_score_function(
-                 scores, topk, num_groups, group_topk, _compute_topk
-             )
+@@ -587,6 +587,9 @@ def topk_routing_with_score_function(
+         else:
+             return torch.topk(scores, k=topk, dim=1)
  
 +    from slime.utils.routing_replay import get_routing_replay_compute_topk
 +    compute_topk = get_routing_replay_compute_topk(compute_topk)
@@ -577,12 +526,12 @@ index 75825cd37..445b3fb84 100644
          if use_pre_softmax:
              scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
 diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
-index a2f3e90bd..b6f732561 100644
+index 16fc9d9af..517944f25 100644
 --- a/megatron/core/transformer/moe/router.py
 +++ b/megatron/core/transformer/moe/router.py
-@@ -207,6 +207,9 @@ class TopKRouter(Router):
-         if self.config.moe_enable_routing_replay:
-             self.router_replay = RouterReplay()
+@@ -201,6 +201,9 @@ class TopKRouter(Router):
+             self.global_tokens_per_expert = None
+             self.ga_steps = None
  
 +        from slime.utils.routing_replay import register_routing_replay
 +        register_routing_replay(self)
@@ -591,10 +540,18 @@ index a2f3e90bd..b6f732561 100644
          """
          Maintain the expert bias in float32.
 diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
-index b0476155a..63f81465d 100755
+index a8f4abfcd..f33f6f05e 100755
 --- a/megatron/core/transformer/multi_token_prediction.py
 +++ b/megatron/core/transformer/multi_token_prediction.py
-@@ -709,17 +709,21 @@ class MultiTokenPredictionLayer(MegatronModule):
+@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Union
+ 
+ import torch
+ from torch import Tensor
++import warnings
+ 
+ from megatron.core import InferenceParams, parallel_state, tensor_parallel
+ from megatron.core.dist_checkpointing.mapping import ShardedStateDict
+@@ -714,17 +715,19 @@ class MultiTokenPredictionLayer(MegatronModule):
              cp_group=self.cp_group,
              packed_seq_params=packed_seq_params,
          )
@@ -618,20 +575,24 @@ index b0476155a..63f81465d 100755
 +        decoder_input = decoder_input.detach()
  
 -        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
-+        hidden_states = make_viewless_tensor(
-+            inp=hidden_states, requires_grad=True, keep_graph=False
-+        )
++        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=False)
  
          return input_ids, position_ids, decoder_input, hidden_states
  
-@@ -821,22 +825,60 @@ class MultiTokenPredictionLayer(MegatronModule):
+@@ -826,6 +829,51 @@ class MultiTokenPredictionLayer(MegatronModule):
          return hidden_states
  
      def _checkpointed_forward(self, forward_func, *args, **kwargs):
-+        """Wrap forward_func with activation checkpointing while only passing tensors."""
++        """Wrap `forward_func` with activation checkpointing while only passing tensors.
++
++        Non-tensor arguments (e.g., configuration objects, None) are captured via closure so
++        that checkpoint implementations never receive them directly, avoiding save_for_backward
++        issues with non-tensor inputs.
++        """
 +
++        # TODO(jiajun): Is there any better implementation here?
 +        positional_specs = []
-+        keyword_specs = []
++        kw_specs = []
 +        tensor_args: List[torch.Tensor] = []
 +
 +        for arg in args:
@@ -643,10 +604,10 @@ index b0476155a..63f81465d 100755
 +
 +        for key, value in kwargs.items():
 +            if torch.is_tensor(value):
-+                keyword_specs.append((key, ('tensor', len(tensor_args))))
++                kw_specs.append((key, ('tensor', len(tensor_args))))
 +                tensor_args.append(value)
 +            else:
-+                keyword_specs.append((key, ('const', value)))
++                kw_specs.append((key, ('const', value)))
 +
 +        def run(*flat_tensor_args):
 +            rebuilt_args = []
@@ -657,7 +618,7 @@ index b0476155a..63f81465d 100755
 +                    rebuilt_args.append(payload)
 +
 +            rebuilt_kwargs = {}
-+            for key, (spec_type, payload) in keyword_specs:
++            for key, (spec_type, payload) in kw_specs:
 +                if spec_type == 'tensor':
 +                    rebuilt_kwargs[key] = flat_tensor_args[payload]
 +                else:
@@ -670,11 +631,7 @@ index b0476155a..63f81465d 100755
          def checkpoint_handler():
              """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
              if self.config.fp8:
-                 from megatron.core.extensions.transformer_engine import te_checkpoint
- 
-                 return te_checkpoint(
--                    forward_func,
-+                    run,
+@@ -836,12 +884,11 @@ class MultiTokenPredictionLayer(MegatronModule):
                      self.config.distribute_saved_activations,
                      tensor_parallel.random.get_cuda_rng_tracker,
                      parallel_state.get_tensor_model_parallel_group(),
@@ -690,10 +647,10 @@ index b0476155a..63f81465d 100755
  
          if self.config.recompute_method == 'uniform':
 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
-index dce438520..de51edaf3 100644
+index e2705bd9f..a0aa109b5 100644
 --- a/megatron/core/transformer/transformer_config.py
 +++ b/megatron/core/transformer/transformer_config.py
-@@ -229,6 +229,9 @@ class TransformerConfig(ModelParallelConfig):
+@@ -210,6 +210,9 @@ class TransformerConfig(ModelParallelConfig):
      attention_output_gate: bool = False
      """Whether to apply output gate to the attention layers."""
  
@@ -704,10 +661,10 @@ index dce438520..de51edaf3 100644
      """Whether to run real-time tests."""
  
 diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
-index 12c248684..227e95862 100644
+index 3ea405770..5a42001b9 100644
 --- a/megatron/core/transformer/transformer_layer.py
 +++ b/megatron/core/transformer/transformer_layer.py
-@@ -224,6 +224,7 @@ class TransformerLayerSubmodules:
+@@ -223,6 +223,7 @@ class TransformerLayerSubmodules:
      input_layernorm: Union[ModuleSpec, type] = IdentityOp
      self_attention: Union[ModuleSpec, type] = IdentityOp
      self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
@@ -715,7 +672,7 @@ index 12c248684..227e95862 100644
  
      pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
      cross_attention: Union[ModuleSpec, type] = IdentityOp
-@@ -232,6 +233,7 @@ class TransformerLayerSubmodules:
+@@ -231,6 +232,7 @@ class TransformerLayerSubmodules:
      pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
      mlp: Union[ModuleSpec, type] = IdentityOp
      mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp
@@ -723,7 +680,7 @@ index 12c248684..227e95862 100644
  
      # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method
      sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
-@@ -311,6 +313,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
+@@ -310,6 +312,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
          # [Module 3: BiasDropoutFusion]
          self.self_attn_bda = build_module(submodules.self_attn_bda)
  
@@ -737,7 +694,7 @@ index 12c248684..227e95862 100644
          # [Module 4: Post SelfAttention] Optional Layernorm after self-attn
          self.pre_cross_attn_layernorm = build_module(
              submodules.pre_cross_attn_layernorm,
-@@ -376,6 +385,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
+@@ -375,6 +384,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
  
          self.is_moe_layer = isinstance(self.mlp, MoELayer)
  
@@ -751,7 +708,7 @@ index 12c248684..227e95862 100644
          self.recompute_input_layernorm = False
          self.recompute_pre_mlp_layernorm = False
          self.recompute_mlp = False
-@@ -615,6 +631,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
+@@ -551,6 +567,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
                  attention_output_with_bias[0]
              )
  
@@ -762,9 +719,9 @@ index 12c248684..227e95862 100644
          # TODO: could we move `bias_dropout_add_exec_handler` itself
          # inside the module provided in the `bias_dropout_add_spec` module?
          nvtx_range_push(suffix="self_attn_bda")
-@@ -794,6 +814,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
-             self.config.inference_fuse_tp_communication
-         )
+@@ -677,6 +697,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
+         else:
+             mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
  
 +        mlp_output, mlp_output_bias = mlp_output_with_bias
 +        mlp_output = self.post_mlp_layernorm(mlp_output)
@@ -774,10 +731,10 @@ index 12c248684..227e95862 100644
              # discard the output of the pre-mlp layernorm and register the recompute
              # as a gradient hook of mlp_output_with_bias[0]
 diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
-index 1af066a82..8c7acadb6 100644
+index b267c8a81..83736acdc 100644
 --- a/megatron/training/arguments.py
 +++ b/megatron/training/arguments.py
-@@ -1388,6 +1388,9 @@ def core_transformer_config_from_args(args, config_class=None):
+@@ -1398,6 +1398,9 @@ def core_transformer_config_from_args(args, config_class=None):
  
      kw_args['inference_sampling_seed'] = args.seed
  
@@ -787,17 +744,21 @@ index 1af066a82..8c7acadb6 100644
      # handle quantization config
      # NOTE: Kitchen arguments are only added to the namespace when
      # Kitchen library is available.
-@@ -1701,6 +1704,8 @@ def _add_network_size_args(parser):
-     group.add_argument('--make-vocab-size-divisible-by', type=int, default=128,
-                        help='Pad the vocab size to be divisible by this value.'
-                        'This is added for computational efficieny reasons.')
+@@ -1764,6 +1767,12 @@ def _add_network_size_args(parser):
+                        action='store_true',
+                        help='If set, use original BERT residula connection '
+                        'ordering.')
++    group.add_argument('--post-self-attn-layernorm', action='store_true',
++                       help='If set, use post self attention layernorm.')
++    group.add_argument('--post-mlp-layernorm', action='store_true',
++                       help='If set, use post MLP layernorm.')
 +    group.add_argument('--use-gated-attention', action='store_true',
 +                       help='If set, use gated attention as in Qwen3Next')
      group.add_argument('--openai-gelu', action='store_true',
                         help='Use OpenAIs GeLU implementation. This option'
                         'should not be used unless for backward compatibility'
 diff --git a/megatron/training/tokenizer/tokenizer.py b/megatron/training/tokenizer/tokenizer.py
-index 17df57dda..260a5f6c8 100644
+index 13b7526ca..6c590f653 100644
 --- a/megatron/training/tokenizer/tokenizer.py
 +++ b/megatron/training/tokenizer/tokenizer.py
 @@ -136,7 +136,7 @@ class _HuggingFaceTokenizer(MegatronLegacyTokenizer):
diff --git a/docker/patch/latest/sglang.patch b/docker/patch/latest/sglang.patch
index 4a13e2f9b4..89f99ab5a5 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 691f06411d..671ac81c48 100644
+index 6fbd1db82..f80ec11bb 100644
 --- a/python/sglang/srt/configs/model_config.py
 +++ b/python/sglang/srt/configs/model_config.py
-@@ -294,6 +294,7 @@ class ModelConfig:
+@@ -274,6 +274,7 @@ class ModelConfig:
  
          if is_draft_model and self.hf_config.architectures[0] in [
              "DeepseekV3ForCausalLM",
@@ -10,20 +10,69 @@ index 691f06411d..671ac81c48 100644
              "GlmMoeDsaForCausalLM",
          ]:
              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 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:
-     kv_data_ptrs: List[int]
-     kv_data_lens: List[int]
-     kv_item_lens: List[int]
-+    aux_buffer_names: List[str]
-     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 67fe82ad6..67f564f04 100644
+--- a/python/sglang/srt/disaggregation/common/conn.py
++++ b/python/sglang/srt/disaggregation/common/conn.py
+@@ -116,10 +116,18 @@ class CommonKVManager(BaseKVManager):
+ 
+         bootstrap_server_url = f"{host}:{self.bootstrap_port}"
+         url = f"http://{bootstrap_server_url}/route"
++        route_attn_tp_rank = self.attn_tp_rank
++        # In prefill CP mode, attention TP rank is flattened to 0, but requests are
++        # still routed by engine rank; register by engine rank to preserve all routes.
++        if (
++            self.disaggregation_mode == DisaggregationMode.PREFILL
++            and self.attn_tp_size == 1
++        ):
++            route_attn_tp_rank = self.kv_args.engine_rank
+         payload = {
+             "role": "Prefill",
+             "attn_tp_size": self.attn_tp_size,
+-            "attn_tp_rank": self.attn_tp_rank,
++            "attn_tp_rank": route_attn_tp_rank,
+             "attn_dp_size": self.attn_dp_size,
+             "attn_dp_rank": self.attn_dp_rank,
+             "pp_size": self.pp_size,
+@@ -333,6 +341,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
+             )
+@@ -357,6 +369,11 @@ class CommonKVReceiver(BaseKVReceiver):
+             # multiple connections in the connection pool and have to send dummy requests to other prefill ranks,
+             # or the KVPoll will never be set correctly
+             self.target_tp_rank = self.target_tp_ranks[0]
++            # For prefill CP mode (decode attention TP=1, prefill attention TP>1),
++            # route bootstrap to all prefill ranks as non-dummy so the serving rank
++            # always receives decode-side metadata.
++            if self.kv_mgr.attn_tp_size == 1 and self.prefill_attn_tp_size > 1:
++                self.target_tp_rank = None
+             self.required_dst_info_num = 1
+             if self.kv_mgr.is_mla_backend:
+                 self.required_prefill_response_num = (
+@@ -610,8 +627,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 f54c882cc2..03832002f0 100644
+index 1d8baf002..1672de78d 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
@@ -34,27 +83,8 @@ index f54c882cc2..03832002f0 100644
  import time
  from collections import deque
  from dataclasses import dataclass
-@@ -42,8 +43,10 @@ from sglang.srt.disaggregation.utils import (
-     MetadataBuffers,
-     ReqToMetadataIdxAllocator,
-     TransferBackend,
-+    apply_prefill_timing_payload,
-     get_kv_class,
-     is_mla_backend,
-+    is_slime_profiling_enabled,
-     kv_to_page_indices,
-     poll_and_all_reduce,
-     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()
+@@ -336,6 +337,16 @@ class DecodePreallocQueue:
          )
-+        kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names()
- 
-         if hasattr(self.token_to_kv_pool, "get_state_buf_infos"):
-             state_data_ptrs, state_data_lens, state_item_lens = (
-@@ -398,6 +402,16 @@ class DecodePreallocQueue:
-                     )
          return kv_manager
  
 +    def release_memory_occupation(self):
@@ -70,7 +100,7 @@ index f54c882cc2..03832002f0 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):
-@@ -525,12 +539,37 @@ class DecodePreallocQueue:
+@@ -440,12 +451,37 @@ class DecodePreallocQueue:
              [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
          )
  
@@ -108,41 +138,10 @@ index f54c882cc2..03832002f0 100644
 +                        self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
              elif poll == KVPoll.WaitingForInput:
                  decode_req.waiting_for_input = True
-                 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.send_metadata(
-                 page_indices, decode_req.metadata_buffer_index, state_indices
-@@ -964,6 +1004,7 @@ class DecodeTransferQueue:
-             output_topk_index,
-             output_hidden_states,
-             output_bootstrap_room,
-+            output_prefill_timing,
-         ) = self.metadata_buffers.get_buf(idx)
- 
-         # Validate bootstrap_room to detect context corruption
-@@ -1025,6 +1066,14 @@ class DecodeTransferQueue:
-                 output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
-             )
- 
-+        # Inject prefill-side PD timing forwarded from the P instance.
-+        # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap,
-+        #          alloc_waiting, transfer_speed, transfer_mb, retry_count]
-+        if is_slime_profiling_enabled():
-+            apply_prefill_timing_payload(
-+                decode_req.req.time_stats, output_prefill_timing
-+            )
-+
-         decode_req.kv_receiver.clear()
-         decode_req.kv_receiver = None
-         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
-             )
+             elif poll == KVPoll.Failed:
+@@ -830,6 +866,13 @@ class DecodeTransferQueue:
+             [decode_req.kv_receiver for decode_req 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.
@@ -154,7 +153,7 @@ index f54c882cc2..03832002f0 100644
          transferred_reqs = []
          indices_to_remove = set()
          for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
-@@ -1111,7 +1167,20 @@ class DecodeTransferQueue:
+@@ -877,7 +920,31 @@ class DecodeTransferQueue:
                  KVPoll.WaitingForInput,
                  KVPoll.Transferring,
              ]:
@@ -172,11 +171,22 @@ index f54c882cc2..03832002f0 100644
 +                        f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
 +                    )
 +                    logger.error(error_message)
-+                    decode_req.kv_receiver.abort()
++                    prepare_abort(
++                        decode_req.req,
++                        error_message,
++                        status_code=HTTPStatus.GATEWAY_TIMEOUT,
++                    )
++                    self.scheduler.stream_output(
++                        [decode_req.req], decode_req.req.return_logprob
++                    )
++                    release_kv_cache(decode_req.req, self.tree_cache, is_insert=False)
++                    indices_to_remove.add(i)
++                    if self.scheduler.enable_metrics:
++                        self.scheduler.metrics_collector.increment_transfer_failed_reqs()
              else:
                  raise ValueError(f"Unexpected poll case: {poll}")
  
-@@ -1132,6 +1201,14 @@ class DecodeTransferQueue:
+@@ -893,6 +960,14 @@ class DecodeTransferQueue:
  
          return transferred_reqs
  
@@ -191,7 +201,7 @@ index f54c882cc2..03832002f0 100644
  
  class SchedulerDisaggregationDecodeMixin:
  
-@@ -1301,7 +1378,15 @@ class SchedulerDisaggregationDecodeMixin:
+@@ -1072,7 +1147,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:
@@ -209,72 +219,30 @@ index f54c882cc2..03832002f0 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 64d97f5c69..4ef08446aa 100644
+index d0d4efd95..a5f06cd67 100644
 --- a/python/sglang/srt/disaggregation/mooncake/conn.py
 +++ b/python/sglang/srt/disaggregation/mooncake/conn.py
-@@ -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.distributed.parallel_state import get_mooncake_transfer_engine
- from sglang.srt.environ import envs
-@@ -276,6 +277,16 @@ class MooncakeKVManager(CommonKVManager):
+@@ -260,6 +260,19 @@ 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)
 +
-     # ------------------------------------------------------------------
-     # 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
- 
--        for i, dst_aux_ptr in enumerate(dst_aux_ptrs):
--            length = prefill_aux_item_lens[i]
--            src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
--            dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index
-+        for _, src_addr, dst_addr, length in iter_aux_transfer_specs(
-+            self.kv_args.aux_buffer_names,
-+            prefill_aux_ptrs,
-+            prefill_aux_item_lens,
-+            dst_aux_ptrs,
-+            prefill_aux_index,
-+            req.dst_aux_index,
-+        ):
-             transfer_blocks.append((src_addr, dst_addr, length))
- 
-         return self._transfer_data(req.mooncake_session_id, transfer_blocks)
-@@ -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
- 
--        for i in range(len(prefill_aux_ptrs)):
--            length = prefill_aux_item_lens[i]
--            src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
-+        for i, src_addr, _, length in iter_aux_transfer_specs(
-+            self.kv_args.aux_buffer_names,
-+            prefill_aux_ptrs,
-+            prefill_aux_item_lens,
-+            dst_aux_ptrs,
-+            prefill_aux_index,
-+            req.dst_aux_index,
-+        ):
-             data = AuxDataCodec.serialize_data_from_buffer(src_addr, length)
- 
-             self.send_aux_data_to_endpoint(
-@@ -1002,13 +1022,13 @@ class MooncakeKVManager(CommonKVManager):
+     def _transfer_data(self, mooncake_session_id, transfer_blocks):
+         if not transfer_blocks:
+             return 0
+@@ -643,13 +656,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."
                  )
@@ -294,22 +262,9 @@ index 64d97f5c69..4ef08446aa 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)
-@@ -1266,12 +1286,6 @@ class MooncakeKVManager(CommonKVManager):
-                         if ret != 0:
-                             with self.session_lock:
-                                 self.session_failures[req.mooncake_session_id] += 1
--                                # Failures should never happen if the session is not dead, if the session fails once, mark it as failed
--                                if self.session_failures[req.mooncake_session_id] >= 1:
--                                    self.failed_sessions.add(req.mooncake_session_id)
--                                    logger.error(
--                                        f"Session {req.mooncake_session_id} failed."
--                                    )
-                             self.record_failure(
-                                 kv_chunk.room,
-                                 f"Failed to send kv chunk of {kv_chunk.room} to "
-@@ -1289,13 +1303,31 @@ class MooncakeKVManager(CommonKVManager):
- 
-                         if kv_chunk.is_last_chunk:
+@@ -880,13 +893,36 @@ class MooncakeKVManager(CommonKVManager):
+ 
+                         if kv_chunk.is_last:
                              if kv_chunk.state_indices is not None:
 -                                self.maybe_send_extra(
 +                                ret = self.maybe_send_extra(
@@ -320,13 +275,18 @@ index 64d97f5c69..4ef08446aa 100644
                                      target_rank_registration_info,
                                  )
 +                                if ret != 0:
-+                                    remote_addr = NetworkAddress(
-+                                        req.endpoint, req.dst_port
-+                                    ).to_host_port_str()
++                                    with self.session_lock:
++                                        self.session_failures[req.mooncake_session_id] += 1
++                                        if self.session_failures[req.mooncake_session_id] >= 1:
++                                            self.failed_sessions.add(
++                                                req.mooncake_session_id
++                                            )
++                                            logger.error(
++                                                f"Session {req.mooncake_session_id} failed."
++                                            )
 +                                    self.record_failure(
 +                                        kv_chunk.room,
-+                                        f"Failed to send extra state chunk of {kv_chunk.room} to "
-+                                        f"{remote_addr}",
++                                        f"Failed to send extra state chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}",
 +                                    )
 +                                    self.update_status(kv_chunk.room, KVPoll.Failed)
 +                                    self.sync_status_to_decode_endpoint(
@@ -334,37 +294,78 @@ index 64d97f5c69..4ef08446aa 100644
 +                                        req.dst_port,
 +                                        req.room,
 +                                        KVPoll.Failed,
-+                                        prefill_unique_rank,
++                                        local_rank,
 +                                    )
 +                                    break
  
                              # Only the last chunk we need to send the aux data
                              ret = self.send_aux(
+@@ -895,6 +931,21 @@ 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
++                                    if (
++                                        self.session_failures[req.mooncake_session_id]
++                                        >= 1
++                                    ):
++                                        self.failed_sessions.add(
++                                            req.mooncake_session_id
++                                        )
++                                        logger.error(
++                                            f"Session {req.mooncake_session_id} failed (send_aux)."
++                                        )
+                             dst_ranks_infos.append(
+                                 (req.endpoint, req.dst_port, req.room)
+                             )
+@@ -977,15 +1028,18 @@ class MooncakeKVManager(CommonKVManager):
+ 
+                 if status == KVPoll.Success:
+                     if bootstrap_room in self.request_status:
+-                        self.prefill_response_tracker[bootstrap_room].add(prefill_rank)
+-                        expected_response_num = (
+-                            self.required_prefill_response_num_table[bootstrap_room]
++                        # 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.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,
 diff --git a/python/sglang/srt/disaggregation/prefill.py b/python/sglang/srt/disaggregation/prefill.py
-index 8eadf81954..c180ce79f3 100644
+index fbc801635..4ea5638cd 100644
 --- a/python/sglang/srt/disaggregation/prefill.py
 +++ b/python/sglang/srt/disaggregation/prefill.py
-@@ -20,6 +20,8 @@ Life cycle of a request in the prefill server
+@@ -20,6 +20,7 @@ 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
- 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()
-         )
-+        kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names()
-         kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
-         kv_args.gpu_id = self.scheduler.gpu_id
- 
-@@ -290,6 +293,11 @@ class PrefillBootstrapQueue:
-             self.scheduler.attn_tp_cpu_group,
+@@ -276,6 +277,12 @@ class PrefillBootstrapQueue:
+             [req.disagg_kv_sender for req in self.queue], self.gloo_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")
 +        )
@@ -373,10 +374,11 @@ index 8eadf81954..c180ce79f3 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
-@@ -297,6 +305,26 @@ class PrefillBootstrapQueue:
+@@ -283,6 +290,27 @@ class PrefillBootstrapQueue:
                      continue
  
              if poll == KVPoll.Bootstrapping:
++                # Check for bootstrap timeout
 +                entry_time = getattr(
 +                    req.time_stats,
 +                    "prefill_bootstrap_queue_entry_time",
@@ -400,7 +402,7 @@ index 8eadf81954..c180ce79f3 100644
                  continue
              elif poll == KVPoll.Failed:
                  error_message = f"Prefill bootstrap failed for request rank={self.tp_rank} {req.rid=} {req.bootstrap_room=}"
-@@ -346,6 +374,15 @@ class PrefillBootstrapQueue:
+@@ -335,6 +363,15 @@ class PrefillBootstrapQueue:
          else:
              return bootstrapped_reqs, failed_reqs
  
@@ -416,34 +418,12 @@ index 8eadf81954..c180ce79f3 100644
  
  class SchedulerDisaggregationPrefillMixin:
      """
-@@ -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.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
-@@ -593,6 +635,11 @@ class SchedulerDisaggregationPrefillMixin:
+@@ -564,6 +601,13 @@ 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")
 +        )
@@ -452,11 +432,18 @@ index 8eadf81954..c180ce79f3 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):
-@@ -618,7 +665,29 @@ class SchedulerDisaggregationPrefillMixin:
+@@ -573,10 +617,35 @@ class SchedulerDisaggregationPrefillMixin:
+                     undone_reqs.append(req)
                      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",
@@ -469,7 +456,7 @@ index 8eadf81954..c180ce79f3 100644
 +                        f"{req.rid=} {req.bootstrap_room=}"
 +                    )
 +                    logger.error(error_message)
-+                    release_kv_cache(req, self.tree_cache)
++                    release_kv_cache(req, self.tree_cache)  # unlock the tree
 +                    prepare_abort(
 +                        req, error_message, status_code=HTTPStatus.GATEWAY_TIMEOUT
 +                    )
@@ -483,218 +470,36 @@ index 8eadf81954..c180ce79f3 100644
              elif poll == KVPoll.Success:  # transfer done
                  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 d7956a6048..0ced278713 100644
---- a/python/sglang/srt/disaggregation/utils.py
-+++ b/python/sglang/srt/disaggregation/utils.py
-@@ -28,6 +28,17 @@ if TYPE_CHECKING:
- # Constants & Enums
- #########################
- FAKE_BOOTSTRAP_HOST = "2.2.2.2"
-+PREFILL_TIMING_AUX_BUFFER_NAME = "prefill_timing"
-+PREFILL_TIMING_DEST_ATTRS = (
-+    ("fwd_prefill_bootstrap_queue_duration", float),
-+    ("fwd_prefill_forward_duration", float),
-+    ("fwd_prefill_transfer_queue_duration", float),
-+    ("fwd_bootstrap_duration", float),
-+    ("fwd_alloc_waiting_duration", float),
-+    ("fwd_transfer_speed_gb_s", float),
-+    ("fwd_transfer_total_mb", float),
-+    ("fwd_prefill_retry_count", int),
-+)
+@@ -743,7 +812,7 @@ class SchedulerDisaggregationPrefillMixin:
  
- 
- class DisaggregationMode(Enum):
-@@ -193,46 +204,35 @@ class MetadataBuffers:
-             self.bootstrap_room = torch.zeros(
-                 (size, 8), dtype=bootstrap_room_dtype, device=device
+         page_indices = kv_to_page_indices(kv_indices, page_size)
+         if len(page_indices) == 0:
+-            logger.info(
++            logger.debug(
+                 f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
              )
-+            # Prefill-side PD timing (8 floats, padded to 16 for RDMA alignment).
-+            # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap,
-+            #          alloc_waiting, transfer_speed, transfer_mb, retry_count]
-+            self.prefill_timing = torch.zeros(
-+                (size, 16), dtype=torch.float32, device=device
-+            )
-+        self.aux_buffers = [
-+            ("output_ids", self.output_ids),
-+            ("cached_tokens", self.cached_tokens),
-+            ("output_token_logprobs_val", self.output_token_logprobs_val),
-+            ("output_token_logprobs_idx", self.output_token_logprobs_idx),
-+            ("output_top_logprobs_val", self.output_top_logprobs_val),
-+            ("output_top_logprobs_idx", self.output_top_logprobs_idx),
-+            ("output_topk_p", self.output_topk_p),
-+            ("output_topk_index", self.output_topk_index),
-+            ("output_hidden_states", self.output_hidden_states),
-+            ("bootstrap_room", self.bootstrap_room),
-+            (PREFILL_TIMING_AUX_BUFFER_NAME, self.prefill_timing),
-+        ]
- 
-     def get_buf_infos(self):
--        ptrs = [
--            self.output_ids.data_ptr(),
--            self.cached_tokens.data_ptr(),
--            self.output_token_logprobs_val.data_ptr(),
--            self.output_token_logprobs_idx.data_ptr(),
--            self.output_top_logprobs_val.data_ptr(),
--            self.output_top_logprobs_idx.data_ptr(),
--            self.output_topk_p.data_ptr(),
--            self.output_topk_index.data_ptr(),
--            self.output_hidden_states.data_ptr(),
--            self.bootstrap_room.data_ptr(),
--        ]
--        data_lens = [
--            self.output_ids.nbytes,
--            self.cached_tokens.nbytes,
--            self.output_token_logprobs_val.nbytes,
--            self.output_token_logprobs_idx.nbytes,
--            self.output_top_logprobs_val.nbytes,
--            self.output_top_logprobs_idx.nbytes,
--            self.output_topk_p.nbytes,
--            self.output_topk_index.nbytes,
--            self.output_hidden_states.nbytes,
--            self.bootstrap_room.nbytes,
--        ]
--        item_lens = [
--            self.output_ids[0].nbytes,
--            self.cached_tokens[0].nbytes,
--            self.output_token_logprobs_val[0].nbytes,
--            self.output_token_logprobs_idx[0].nbytes,
--            self.output_top_logprobs_val[0].nbytes,
--            self.output_top_logprobs_idx[0].nbytes,
--            self.output_topk_p[0].nbytes,
--            self.output_topk_index[0].nbytes,
--            self.output_hidden_states[0].nbytes,
--            self.bootstrap_room[0].nbytes,
--        ]
-+        ptrs = [buffer.data_ptr() for _, buffer in self.aux_buffers]
-+        data_lens = [buffer.nbytes for _, buffer in self.aux_buffers]
-+        item_lens = [buffer[0].nbytes for _, buffer in self.aux_buffers]
-         return ptrs, data_lens, item_lens
- 
-+    def get_aux_buffer_names(self):
-+        return [name for name, _ in self.aux_buffers]
-+
-     def get_buf(self, idx: int):
-         return (
-             self.output_ids[idx],
-@@ -245,8 +245,12 @@ class MetadataBuffers:
-             self.output_topk_index[idx],
-             self.output_hidden_states[idx],
-             self.bootstrap_room[idx],
-+            self.prefill_timing[idx],
-         )
- 
-+    def clear_profiling_buf(self, idx: int):
-+        self.prefill_timing[idx].zero_()
-+
-     def set_buf(self, req: Req):
- 
-         self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
-@@ -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
-         )
-+        # Pack prefill-side PD timing durations for transfer to decode instance.
-+        # Note: set_buf is called at the START of the last KV chunk send, so
-+        # completion_time and prefill_transfer_queue_entry_time are not yet set.
-+        # We use time.perf_counter() as the "forward just completed" timestamp.
-+        import time
-+
-+        ts = req.time_stats
-+        timing = self.prefill_timing[req.metadata_buffer_index]
-+        self.clear_profiling_buf(req.metadata_buffer_index)
-+        if not is_slime_profiling_enabled():
-+            return
-+        for idx, value in enumerate(
-+            build_prefill_timing_payload(ts, now=time.perf_counter())
-+        ):
-+            if value > 0:
-+                timing[idx] = value
-+
-+
-+def is_slime_profiling_enabled() -> bool:
-+    return envs.SLIME_ENABLE_PROFILING.get()
-+
-+
-+def build_prefill_timing_payload(time_stats, now: float) -> tuple[float, ...]:
-+    bootstrap_queue_duration = 0.0
-+    if (
-+        time_stats.prefill_bootstrap_queue_entry_time > 0
-+        and time_stats.wait_queue_entry_time > 0
-+    ):
-+        bootstrap_queue_duration = (
-+            time_stats.wait_queue_entry_time
-+            - time_stats.prefill_bootstrap_queue_entry_time
-+        )
-+
-+    prefill_forward_duration = (
-+        now - time_stats.forward_entry_time
-+        if time_stats.forward_entry_time > 0
-+        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, 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)),
-+    )
-+
-+
-+def apply_prefill_timing_payload(time_stats, timing) -> None:
-+    for value, (attr_name, caster) in zip(
-+        timing[: len(PREFILL_TIMING_DEST_ATTRS)].tolist(),
-+        PREFILL_TIMING_DEST_ATTRS,
-+    ):
-+        if value > 0:
-+            setattr(time_stats, attr_name, caster(value))
-+
-+
-+def iter_aux_transfer_specs(
-+    aux_buffer_names: list[str],
-+    prefill_aux_ptrs: list[int],
-+    prefill_aux_item_lens: list[int],
-+    dst_aux_ptrs: list[int],
-+    prefill_aux_index: int,
-+    dst_aux_index: int,
-+):
-+    profiling_enabled = is_slime_profiling_enabled()
-+    for i, (buffer_name, dst_aux_ptr) in enumerate(zip(aux_buffer_names, dst_aux_ptrs)):
-+        if not profiling_enabled and buffer_name == PREFILL_TIMING_AUX_BUFFER_NAME:
-+            continue
-+        length = prefill_aux_item_lens[i]
-+        if length <= 0:
-+            continue
-+        src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
-+        dst_addr = dst_aux_ptr + length * dst_aux_index
-+        yield i, src_addr, dst_addr, length
+             return
+diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py
+index 8f1069c00..e47589295 100644
+--- a/python/sglang/srt/distributed/parallel_state.py
++++ b/python/sglang/srt/distributed/parallel_state.py
+@@ -1999,7 +1999,10 @@ def get_tensor_model_parallel_world_size():
+ 
+ def get_tensor_model_parallel_rank():
+     """Return my rank for the tensor model parallel group."""
+-    return get_tp_group().rank_in_group
++    try:
++        return get_tp_group().rank_in_group
++    except Exception:
++        return 0
  
  
- #########################
+ # ATTN_TP
 diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
-index d864e4abaa..3a000a80f2 100644
+index 0ed5a1b44..67e33c650 100644
 --- a/python/sglang/srt/entrypoints/engine.py
 +++ b/python/sglang/srt/entrypoints/engine.py
-@@ -69,6 +69,7 @@ from sglang.srt.managers.io_struct import (
+@@ -52,6 +52,7 @@ from sglang.srt.managers.io_struct import (
      LoadLoRAAdapterReqInput,
      MultimodalDataInputFormat,
      OpenSessionReqInput,
@@ -702,7 +507,7 @@ index d864e4abaa..3a000a80f2 100644
      ReleaseMemoryOccupationReqInput,
      ResumeMemoryOccupationReqInput,
      RpcReqInput,
-@@ -957,6 +958,24 @@ class Engine(EngineScoreMixin, EngineBase):
+@@ -641,6 +642,24 @@ class Engine(EngineBase):
              self.tokenizer_manager.update_weights_from_ipc(obj, None)
          )
  
@@ -728,10 +533,10 @@ index d864e4abaa..3a000a80f2 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 6978e0c062..80dc159e8f 100644
+index 1d6816c01..402b42e05 100644
 --- a/python/sglang/srt/entrypoints/http_server.py
 +++ b/python/sglang/srt/entrypoints/http_server.py
-@@ -127,6 +127,7 @@ from sglang.srt.managers.io_struct import (
+@@ -115,6 +115,7 @@ from sglang.srt.managers.io_struct import (
      OpenSessionReqInput,
      ParseFunctionCallReq,
      PauseGenerationReqInput,
@@ -739,7 +544,7 @@ index 6978e0c062..80dc159e8f 100644
      ProfileReqInput,
      ReleaseMemoryOccupationReqInput,
      ResumeMemoryOccupationReqInput,
-@@ -582,10 +583,8 @@ async def model_info():
+@@ -574,10 +575,8 @@ async def model_info():
  @app.get("/weight_version")
  async def weight_version():
      """Get the current weight version."""
@@ -752,7 +557,7 @@ index 6978e0c062..80dc159e8f 100644
  
  
  @app.get("/get_server_info")
-@@ -602,9 +601,19 @@ async def get_server_info():
+@@ -594,9 +593,19 @@ async def get_server_info():
  async def server_info():
      """Get the server information."""
      # Returns internal states per DP.
@@ -775,7 +580,7 @@ index 6978e0c062..80dc159e8f 100644
  
      # This field is not serializable.
      if hasattr(_global_state.tokenizer_manager.server_args, "model_config"):
-@@ -1121,6 +1130,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re
+@@ -1084,6 +1093,23 @@ async def update_weights_from_ipc(obj: UpdateWeightsFromIPCReqInput, request: Re
          return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
  
  
@@ -799,20 +604,28 @@ index 6978e0c062..80dc159e8f 100644
  @app.post("/update_weight_version")
  @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 dfc5507de0..be9501b05a 100644
---- a/python/sglang/srt/environ.py
-+++ b/python/sglang/srt/environ.py
-@@ -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)
-     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/index_buf_accessor.py b/python/sglang/srt/layers/attention/nsa/index_buf_accessor.py
+index 1cdf65b91..4783cd18f 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,
 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
-index 02ef4e2440..fd5a43cce8 100644
+index ca54a931b..6c102a251 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 @@
@@ -823,44 +636,23 @@ index 02ef4e2440..fd5a43cce8 100644
  from abc import ABC, abstractmethod
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
  
-@@ -213,14 +214,31 @@ class Indexer(MultiPlatformOp):
-             prefix=add_prefix("weights_proj", prefix),
-         )
-         self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
-+        server_args = get_global_server_args()
-+        disable_flag = server_args.disable_indexer_rope_neox_style
-+        env_raw = os.environ.get("INDEXER_ROPE_NEOX_STYLE", None)
-+        if env_raw is not None:
-+            env_value = env_raw == "1"
-+            if disable_flag and env_value:
-+                raise ValueError(
-+                    "Conflict: --disable-indexer-rope-neox-style is set but "
-+                    "INDEXER_ROPE_NEOX_STYLE='1'. "
-+                    "Please remove one or make them consistent."
-+                )
-+            resolved_neox_style = env_value
-+        elif disable_flag:
-+            resolved_neox_style = False
-+        else:
-+            resolved_neox_style = is_neox_style
-+
-         self.rotary_emb = get_rope_wrapper(
-             rope_head_dim,
-             rotary_dim=rope_head_dim,
+@@ -207,7 +208,11 @@ class Indexer(MultiPlatformOp):
              max_position=max_position_embeddings,
              base=rope_theta,  # type: ignore
              rope_scaling=rope_scaling,
 -            is_neox_style=is_neox_style,
--            device=get_global_server_args().device,
-+            is_neox_style=resolved_neox_style,
-+            device=server_args.device,
++            is_neox_style=(
++                os.environ.get("INDEXER_ROPE_NEOX_STYLE", "1") == "1"
++                if os.environ.get("INDEXER_ROPE_NEOX_STYLE", None)
++                else is_neox_style
++            ),
+             device=get_global_server_args().device,
          )
          self.block_size = block_size
-         self.scale_fmt = scale_fmt
-@@ -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)
+@@ -244,6 +249,11 @@ class Indexer(MultiPlatformOp):
+             x = x.to(self.weights_proj.weight.dtype)
+         weights, _ = self.weights_proj(x)
+         weights = weights.float()
 +        if weights.shape[1] < q_scale.shape[1]:
 +            assert q_scale.shape[1] % weights.shape[1] == 0
 +            weights = weights.repeat_interleave(
@@ -869,7 +661,7 @@ index 02ef4e2440..fd5a43cce8 100644
          weights = weights * self.n_heads**-0.5
          weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
          return weights
-@@ -1078,6 +1101,9 @@ class Indexer(MultiPlatformOp):
+@@ -982,15 +992,26 @@ class Indexer(MultiPlatformOp):
              query, key = self._get_q_k_bf16(
                  q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
              )
@@ -878,8 +670,15 @@ index 02ef4e2440..fd5a43cce8 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):
-                 self._store_index_k_cache(
-@@ -1092,6 +1118,9 @@ class Indexer(MultiPlatformOp):
+                 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:
              query, key = self._get_q_k_bf16(
                  q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
              )
@@ -889,11 +688,539 @@ index 02ef4e2440..fd5a43cce8 100644
  
              if enable_dual_stream:
                  current_stream = torch.cuda.current_stream()
+diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py
+index 00ef96f9b..9885fcd0d 100644
+--- a/python/sglang/srt/layers/attention/nsa/utils.py
++++ b/python/sglang/srt/layers/attention/nsa/utils.py
+@@ -54,7 +54,12 @@ def can_nsa_prefill_cp_round_robin_split(forward_batch: "ForwardBatch"):
+         return False
+     cp_size = get_attention_cp_size()
+     seq_len = sum(forward_batch.extend_seq_lens_cpu)
+-    return is_nsa_prefill_cp_round_robin_split() and seq_len > 0 and cp_size > 1
++    return (
++        is_nsa_prefill_cp_round_robin_split()
++        and seq_len >= cp_size
++        and seq_len % cp_size == 0
++        and cp_size > 1
++    )
+ 
+ 
+ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
+@@ -91,20 +96,29 @@ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]):
+ def cal_padded_tokens(forward_batch: "ForwardBatch"):
+     # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch,
+     # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode.
+-    global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
++    if forward_batch.global_num_tokens_cpu is None:
++        # PD prefill CP+PP path can bypass MLP-sync metadata. Reconstruct a single-rank
++        # global token view from the local token count for NSA padding logic.
++        local_tokens = forward_batch.num_token_non_padded_cpu
++        if local_tokens is None:
++            local_tokens = len(forward_batch.input_ids)
++        global_num_tokens = [local_tokens * get_attention_cp_size()]
++    else:
++        global_num_tokens = forward_batch.global_num_tokens_cpu.copy()
+     sync_group_size = len(global_num_tokens)
+     attn_cp_size = get_attention_cp_size()
+     for i in range(sync_group_size):
+         global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_cp_size)
+-    dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
+-        forward_batch.is_extend_in_batch, global_num_tokens
+-    )
+-    if dp_padding_mode.is_max_len():
+-        tokens = max(global_num_tokens)
+-    elif len(global_num_tokens) > 1:
+-        tokens = global_num_tokens[get_attention_dp_rank()]
+-    else:
++    if len(global_num_tokens) == 1:
+         tokens = global_num_tokens[0]
++    else:
++        dp_padding_mode = DpPaddingMode.get_dp_padding_mode(
++            forward_batch.is_extend_in_batch, global_num_tokens
++        )
++        if dp_padding_mode.is_max_len():
++            tokens = max(global_num_tokens)
++        else:
++            tokens = global_num_tokens[get_attention_dp_rank()]
+     if can_nsa_prefill_cp_round_robin_split(forward_batch):
+         tokens = ceil_div(tokens, attn_cp_size)
+     return tokens
+@@ -152,10 +166,9 @@ class NSAContextParallelMetadata:
+ 
+ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
+     if is_nsa_prefill_cp_round_robin_split():
++        if seq_len < cp_size or seq_len % cp_size != 0:
++            return False
+         cur_cp_seq_len = seq_len // cp_size
+-        assert (
+-            seq_len % cp_size == 0
+-        ), f"seq_len {seq_len} is not divisible by cp_size {cp_size} when nsa_prefill_cp_mode is round-robin-split"
+     else:
+         # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
+         # Note: (self.cp_size * 2) To achieve load balancing for seq computation,
+@@ -175,10 +188,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
+ 
+ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
+     if is_nsa_prefill_cp_round_robin_split():
+-        cp_size = get_attention_cp_size()
+-        assert (
+-            input_.shape[0] % cp_size == 0
+-        ), f"Expect input shape 0 can divided by cp size, but got input shape {input_.shape}, cp size {cp_size}"
+         return nsa_cp_round_robin_split_data(input_)
+ 
+     input_list = list(
+@@ -192,11 +201,6 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
+ 
+ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
+     if is_nsa_prefill_cp_round_robin_split():
+-        cp_size = get_attention_cp_size()
+-        assert positions.shape[0] % cp_size == 0, (
+-            f"Expect positions shape 0 can divided by cp size, but got positions shape {positions.shape}, "
+-            f"cp size {cp_size}"
+-        )
+         return nsa_cp_round_robin_split_data(positions)
+ 
+     position_id_list = list(
+diff --git a/python/sglang/srt/layers/communicator_nsa_cp.py b/python/sglang/srt/layers/communicator_nsa_cp.py
+index 296d14568..f4606a769 100644
+--- a/python/sglang/srt/layers/communicator_nsa_cp.py
++++ b/python/sglang/srt/layers/communicator_nsa_cp.py
+@@ -34,7 +34,6 @@ from sglang.srt.layers.communicator import (
+ from sglang.srt.layers.dp_attention import (
+     attn_cp_all_gather_into_tensor,
+     attn_cp_reduce_scatter_tensor,
+-    get_local_dp_buffer,
+ )
+ from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+ 
+@@ -153,9 +152,23 @@ class NSACPCommunicateWithAllReduceAndLayerNormFn(
+         # for decode: attn tp full -> full
+         if nsa_use_prefill_cp(forward_batch):
+             assert context.attn_dp_size == 1
+-            hidden_states, local_hidden_states = (
+-                get_local_dp_buffer(),
+-                hidden_states,
++            local_hidden_states = hidden_states
++            total_tokens = (
++                sum(forward_batch.extend_seq_lens_cpu)
++                if forward_batch.extend_seq_lens_cpu is not None
++                else local_hidden_states.shape[0] * context.attn_cp_size
++            )
++            max_len = (total_tokens + context.attn_cp_size - 1) // context.attn_cp_size
++            if local_hidden_states.shape[0] < max_len:
++                pad = local_hidden_states.new_zeros(
++                    (
++                        max_len - local_hidden_states.shape[0],
++                        local_hidden_states.shape[1],
++                    )
++                )
++                local_hidden_states = torch.cat([local_hidden_states, pad], dim=0)
++            hidden_states = local_hidden_states.new_empty(
++                (max_len * context.attn_cp_size, local_hidden_states.shape[1])
+             )
+             attn_cp_all_gather_into_tensor(
+                 hidden_states,
+diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py
+index 5bf5aa0c8..e52f39fd8 100644
+--- a/python/sglang/srt/layers/dp_attention.py
++++ b/python/sglang/srt/layers/dp_attention.py
+@@ -90,11 +90,11 @@ class _DpGatheredBufferWrapper:
+     _hidden_size: int
+     _dtype: torch.dtype
+     _device: torch.device
+-    _global_dp_buffer_len: int
+-    _local_dp_buffer_len: int
+-    _dp_max_padding: bool
+-    _global_num_tokens: Optional[List[int]]
+-    _is_extend_in_batch: bool
++    _global_dp_buffer_len: int = 0
++    _local_dp_buffer_len: int = 0
++    _dp_max_padding: bool = False
++    _global_num_tokens: Optional[List[int]] = None
++    _is_extend_in_batch: bool = False
+ 
+     @classmethod
+     def set_metadata(cls, hidden_size: int, dtype: torch.dtype, device: torch.device):
+diff --git a/python/sglang/srt/layers/layernorm.py b/python/sglang/srt/layers/layernorm.py
+index 39832c45a..c3d3c6b54 100644
+--- a/python/sglang/srt/layers/layernorm.py
++++ b/python/sglang/srt/layers/layernorm.py
+@@ -93,20 +93,17 @@ class RMSNorm(MultiPlatformOp):
+         eps: float = 1e-6,
+         var_hidden_size: Optional[int] = None,
+         cast_x_before_out_mul: bool = False,
+-        fp32_residual: bool = False,
++        fp32_residual: bool = True,
+         has_weight: bool = True,
+-        weight_dtype: Optional = None,
+-        override_orig_dtype: Optional = None,
+     ) -> None:
+         super().__init__()
+         self.has_weight = has_weight
+         self.cast_x_before_out_mul = cast_x_before_out_mul
+         self.fp32_residual = fp32_residual
+-        self.override_orig_dtype = override_orig_dtype
+         if self.has_weight:
+-            self.weight = nn.Parameter(torch.ones(hidden_size, dtype=weight_dtype))
++            self.weight = nn.Parameter(torch.ones(hidden_size))
+         else:
+-            self.weight = torch.ones(hidden_size, dtype=weight_dtype)
++            self.weight = torch.ones(hidden_size)
+         self.variance_epsilon = eps
+         self.hidden_size = hidden_size
+         self.variance_size_override = (
+@@ -219,16 +216,19 @@ class RMSNorm(MultiPlatformOp):
+     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+         if not x.is_contiguous():
+             x = x.contiguous()
+-        orig_dtype = self.override_orig_dtype or x.dtype
++        orig_dtype = x.dtype
++
++        if residual is not None and not self.fp32_residual:
++            x = x + residual
++            if post_residual_addition is not None:
++                x = x + post_residual_addition
++            residual = x.clone()
+         x = x.to(torch.float32)
+-        if residual is not None:
++        if residual is not None and self.fp32_residual:
+             x = x + residual.to(torch.float32)
+             if post_residual_addition is not None:
+                 x = x + post_residual_addition.to(torch.float32)
+-            if self.fp32_residual:
+-                residual = x.clone()
+-            else:
+-                residual = x.to(orig_dtype)
++            residual = x.to(orig_dtype)
+ 
+         hidden_size = x.shape[-1]
+         if hidden_size != self.hidden_size:
+@@ -314,7 +314,7 @@ class RMSNorm(MultiPlatformOp):
+ 
+             if get_tensor_model_parallel_world_size() > 1:
+                 if post_residual_addition is not None:
+-                    residual = residual + post_residual_addition
++                    x = x + post_residual_addition
+                 fused_result = flashinfer_allreduce_residual_rmsnorm(
+                     input_tensor=x,
+                     residual=residual,
+diff --git a/python/sglang/srt/layers/logits_processor.py b/python/sglang/srt/layers/logits_processor.py
+index aff05bf42..130359232 100644
+--- a/python/sglang/srt/layers/logits_processor.py
++++ b/python/sglang/srt/layers/logits_processor.py
+@@ -872,11 +872,6 @@ class LogitsProcessor(nn.Module):
+                     None,  # bias
+                     True,  # is_vnni
+                 )
+-            elif get_global_server_args().rl_on_policy_target is not None:
+-                # Due to tie-weight, we may not be able to change lm_head's weight dtype
+-                logits = torch.matmul(
+-                    hidden_states.bfloat16(), lm_head.weight.T.bfloat16()
+-                )
+             else:
+                 logits = torch.matmul(
+                     hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
+diff --git a/python/sglang/srt/layers/moe/ep_moe/deepep_bf16_kernels.py b/python/sglang/srt/layers/moe/ep_moe/deepep_bf16_kernels.py
+new file mode 100644
+index 000000000..8d3d0f92e
+--- /dev/null
++++ b/python/sglang/srt/layers/moe/ep_moe/deepep_bf16_kernels.py
+@@ -0,0 +1,146 @@
++"""Fused Triton kernels for DeepEP BF16 low-latency MoE decode.
++
++Replaces the naive activation + masking pipeline (5+ CUDA kernels for silu+mul
++and arange+comparison+masked_fill+copy) with a single Triton elementwise kernel,
++while keeping cuBLAS batched GEMM for the matrix multiplies.
++
++Pipeline: bmm → fused_act_mul_masked (in-place) → bmm(out=hidden)
++  (3 ops total: 2 cuBLAS + 1 Triton, vs original 7-8 separate CUDA kernels)
++"""
++
++import torch
++import triton
++import triton.language as tl
++
++
++@triton.jit
++def _silu_mul_masked_kernel(
++    gate_up_ptr,
++    masked_m_ptr,
++    M,
++    N,
++    stride_ge,
++    stride_gm,
++    stride_gn,
++    BLOCK: tl.constexpr,
++):
++    """Fused SiLU(gate) * up with per-expert masking, written in-place.
++
++    gate_up: [E, M, 2*N] — first N cols are gate, last N cols are up.
++    Writes SiLU(gate)*up to gate_up[:,:,:N] in-place.
++    Rows m >= masked_m[e] are zeroed.
++    """
++    expert_id = tl.program_id(1)
++    pid = tl.program_id(0)
++
++    expert_valid_m = tl.load(masked_m_ptr + expert_id)
++
++    offs = pid * BLOCK + tl.arange(0, BLOCK)
++    total = M * N
++    mask = offs < total
++
++    m = offs // N
++    n = offs % N
++
++    gate_base = gate_up_ptr + expert_id * stride_ge
++
++    gate_val = tl.load(gate_base + m * stride_gm + n * stride_gn, mask=mask, other=0.0)
++    up_val = tl.load(
++        gate_base + m * stride_gm + (n + N) * stride_gn, mask=mask, other=0.0
++    )
++
++    gate_f32 = gate_val.to(tl.float32)
++    result = (gate_f32 * tl.sigmoid(gate_f32)) * up_val.to(tl.float32)
++
++    # Zero invalid rows
++    valid = m < expert_valid_m
++    result = tl.where(valid, result, 0.0)
++
++    tl.store(
++        gate_base + m * stride_gm + n * stride_gn,
++        result.to(gate_up_ptr.dtype.element_ty),
++        mask=mask,
++    )
++
++
++@triton.jit
++def _gelu_mul_masked_kernel(
++    gate_up_ptr,
++    masked_m_ptr,
++    M,
++    N,
++    stride_ge,
++    stride_gm,
++    stride_gn,
++    BLOCK: tl.constexpr,
++):
++    """Fused GELU(gate) * up with per-expert masking, written in-place."""
++    expert_id = tl.program_id(1)
++    pid = tl.program_id(0)
++
++    expert_valid_m = tl.load(masked_m_ptr + expert_id)
++
++    offs = pid * BLOCK + tl.arange(0, BLOCK)
++    total = M * N
++    mask = offs < total
++
++    m = offs // N
++    n = offs % N
++
++    gate_base = gate_up_ptr + expert_id * stride_ge
++
++    gate_val = tl.load(gate_base + m * stride_gm + n * stride_gn, mask=mask, other=0.0)
++    up_val = tl.load(
++        gate_base + m * stride_gm + (n + N) * stride_gn, mask=mask, other=0.0
++    )
++
++    g = gate_val.to(tl.float32)
++    kAlpha = 0.7978845608028654
++    gate_act = 0.5 * g * (1.0 + tl.math.tanh(kAlpha * (g + 0.044715 * g * g * g)))
++    result = gate_act * up_val.to(tl.float32)
++
++    valid = m < expert_valid_m
++    result = tl.where(valid, result, 0.0)
++
++    tl.store(
++        gate_base + m * stride_gm + n * stride_gn,
++        result.to(gate_up_ptr.dtype.element_ty),
++        mask=mask,
++    )
++
++
++def fused_act_mul_masked_inplace(
++    gate_up: torch.Tensor,
++    intermediate_size: int,
++    masked_m: torch.Tensor,
++    use_gelu: bool = False,
++) -> None:
++    """Fused activation + multiply + masking, written in-place to gate_up[:,:,:I].
++
++    After this call, gate_up[:, :, :intermediate_size] contains the masked
++    activated intermediate, suitable for the down projection GEMM.
++
++    Args:
++        gate_up: [E, M, 2*I] output of bmm(tokens, w13.T), modified in-place
++        intermediate_size: I
++        masked_m: [E] per-expert valid token count
++        use_gelu: use GELU instead of SiLU
++    """
++    E, M, _ = gate_up.shape
++    N = intermediate_size
++
++    total = M * N
++    BLOCK = 1024
++    grid = (triton.cdiv(total, BLOCK), E)
++
++    kernel = _gelu_mul_masked_kernel if use_gelu else _silu_mul_masked_kernel
++    kernel[grid](
++        gate_up,
++        masked_m,
++        M,
++        N,
++        gate_up.stride(0),
++        gate_up.stride(1),
++        gate_up.stride(2),
++        BLOCK=BLOCK,
++    )
+diff --git a/python/sglang/srt/layers/moe/ep_moe/layer.py b/python/sglang/srt/layers/moe/ep_moe/layer.py
+index ebcc696ec..3b527021a 100644
+--- a/python/sglang/srt/layers/moe/ep_moe/layer.py
++++ b/python/sglang/srt/layers/moe/ep_moe/layer.py
+@@ -132,11 +132,12 @@ class DeepEPMoE(FusedMoE):
+             and not _is_npu
+             and not (
+                 get_moe_runner_backend().is_flashinfer_cutedsl()
++                and self.quant_config is not None
+                 and self.quant_config.get_name() == "modelopt_fp4"
+             )
++            and (self.use_fp8_w8a8 or self.use_w4afp8)
+         ):
+-            # NPU supports low_latency deepep without deepgemm
+-            # FP4 quantization with flashinfer_cutedsl also supports low_latency deepep without deepgemm
++            # BF16 models don't need deep_gemm; they use per-expert torch.mm
+             assert (
+                 deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+             ), f"DeepEP {self.deepep_mode} mode requires deep_gemm"
+@@ -154,6 +155,10 @@ class DeepEPMoE(FusedMoE):
+             # the last one is invalid rank_id
+             self.expert_mask[:-1] = 1
+ 
++        # Set bf16_weights flag on dispatcher so dispatch skips FP8 quantization
++        if not self.use_fp8_w8a8 and not self.use_w4afp8:
++            self.dispatcher.set_quant_config({"bf16_weights": True})
++
+     def forward(
+         self,
+         hidden_states: torch.Tensor,
+@@ -228,6 +233,8 @@ class DeepEPMoE(FusedMoE):
+         elif DispatchOutputChecker.format_is_deepep_normal(dispatch_output):
+             if self.use_w4afp8:
+                 output = self.forward_cutlass_w4afp8(dispatch_output)
++            elif not self.use_fp8_w8a8:
++                output = self.forward_bf16_normal(dispatch_output)
+             else:
+                 assert False, "forward_deepgemm_contiguous is deprecated"
+         elif DispatchOutputChecker.format_is_deepep_ll(dispatch_output):
+@@ -238,6 +245,8 @@ class DeepEPMoE(FusedMoE):
+                 output = self.forward_flashinfer_cutedsl(dispatch_output)
+             elif self.use_w4afp8:
+                 output = self.forward_cutlass_w4afp8_masked(dispatch_output)
++            elif not self.use_fp8_w8a8:
++                output = self.forward_bf16_ll(dispatch_output)
+             else:
+                 assert False, "forward_deepgemm_masked is deprecated"
+ 
+@@ -341,6 +350,71 @@ class DeepEPMoE(FusedMoE):
+             dispatch_output=dispatch_output,
+         )
+ 
++    def forward_bf16_normal(
++        self,
++        dispatch_output: DeepEPNormalDispatchOutput,
++    ) -> torch.Tensor:
++        from sglang.srt.layers.moe.fused_moe_triton.fused_moe import fused_experts
++
++        hidden_states = dispatch_output.hidden_states
++        topk_ids = dispatch_output.topk_ids
++        topk_weights = dispatch_output.topk_weights
++
++        if hidden_states.shape[0] == 0:
++            return hidden_states
++
++        # topk_ids uses local expert IDs (0..num_local_experts-1), -1 for remote.
++        # fused_experts handles -1 via moe_align_block_size filtering.
++        return fused_experts(
++            hidden_states=hidden_states,
++            w1=self.w13_weight,
++            w2=self.w2_weight,
++            topk_output=(topk_weights, topk_ids, None),
++            moe_runner_config=self.moe_runner_config,
++        )
++
++    def forward_bf16_ll(
++        self,
++        dispatch_output: DeepEPLLDispatchOutput,
++    ) -> torch.Tensor:
++        from sglang.srt.layers.moe.ep_moe.deepep_bf16_kernels import (
++            fused_act_mul_masked_inplace,
++        )
++
++        hidden_states = dispatch_output.hidden_states
++        masked_m = dispatch_output.masked_m
++        expected_m = dispatch_output.expected_m
++
++        _, max_tokens, _ = hidden_states.shape
++        if masked_m.numel() == 0 or max_tokens == 0:
++            return hidden_states
++
++        expected_m = min(expected_m, max_tokens)
++        if expected_m <= 0:
++            return hidden_states
++
++        tokens = hidden_states[:, :expected_m, :]
++
++        # 1. Gate+Up GEMM (cuBLAS batched GEMM)
++        gate_up = torch.bmm(tokens, self.w13_weight.transpose(1, 2))
++
++        # 2. Fused SiLU(gate)*up + masking in-place (1 Triton kernel replaces 6 ops)
++        fused_act_mul_masked_inplace(
++            gate_up,
++            self.intermediate_size_per_partition,
++            masked_m,
++            use_gelu=(self.moe_runner_config.activation == "gelu"),
++        )
++
++        # 3. Down GEMM into hidden_states (cuBLAS, non-contiguous input is OK)
++        torch.bmm(
++            gate_up[:, :, : self.intermediate_size_per_partition],
++            self.w2_weight.transpose(1, 2),
++            out=hidden_states[:, :expected_m, :],
++        )
++
++        return hidden_states
++
+     def forward_npu(
+         self,
+         dispatch_output: Union[DeepEPNormalDispatchOutput, DeepEPLLDispatchOutput],
+diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+index ebdbb42c6..714ffbe0e 100644
+--- a/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
++++ b/python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py
+@@ -14,6 +14,7 @@ import torch.nn.functional as F
+ import triton.language as tl
+ 
+ from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
++from sglang.srt.server_args import get_global_server_args
+ from sglang.srt.utils import (
+     cpu_has_amx_support,
+     get_bool_env_var,
+@@ -617,7 +618,10 @@ def fused_experts_impl(
+                 ).squeeze(dim=1)
+             else:
+                 # According to micro benchmark results, torch.compile can get better performance for small token.
+-                if tokens_in_chunk <= 32:
++                if (
++                    not get_global_server_args().enable_deterministic_inference
++                    and tokens_in_chunk <= 32
++                ):
+                     moe_sum_reduce_torch_compile(
+                         intermediate_cache3.view(*intermediate_cache3.shape),
+                         out_hidden_states[begin_chunk_idx:end_chunk_idx],
 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 72483f4ea6..2e1148d189 100644
+index de8a07ab3..5c9f4813a 100644
 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
 +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py
-@@ -702,6 +702,7 @@ class FusedMoE(torch.nn.Module):
+@@ -697,6 +697,7 @@ class FusedMoE(torch.nn.Module):
                      "CompressedTensorsWNA16TritonMoE",
                  ]
              )
@@ -901,7 +1228,7 @@ index 72483f4ea6..2e1148d189 100644
              else loaded_weight
          )
  
-@@ -921,6 +922,7 @@ class FusedMoE(torch.nn.Module):
+@@ -916,6 +917,7 @@ class FusedMoE(torch.nn.Module):
                      "CompressedTensorsWNA16TritonMoE",
                  ]
              )
@@ -910,7 +1237,7 @@ index 72483f4ea6..2e1148d189 100644
          )
  
 diff --git a/python/sglang/srt/layers/moe/routed_experts_capturer.py b/python/sglang/srt/layers/moe/routed_experts_capturer.py
-index 00bd687555..12d5577af2 100644
+index 00bd68755..12d5577af 100644
 --- a/python/sglang/srt/layers/moe/routed_experts_capturer.py
 +++ b/python/sglang/srt/layers/moe/routed_experts_capturer.py
 @@ -8,10 +8,15 @@ import torch
@@ -970,11 +1297,79 @@ index 00bd687555..12d5577af2 100644
          self.device_cache.capture_fwd_routed_experts(layer_id, topk_ids)
  
      def get_routed_experts(
+diff --git a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+index 8539639d5..e7f5d1565 100644
+--- a/python/sglang/srt/layers/moe/token_dispatcher/deepep.py
++++ b/python/sglang/srt/layers/moe/token_dispatcher/deepep.py
+@@ -388,6 +388,7 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
+             deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+             and not get_moe_runner_backend().is_cutlass()
+             and not envs.SGLANG_DEEPEP_BF16_DISPATCH.get()
++            and not self.quant_config.get("bf16_weights", False)
+         ):
+             # TODO hard code 128 block quant,use fp8 communication
+             hidden_states = sglang_per_token_group_quant_fp8(
+@@ -466,7 +467,12 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
+             previous_event=previous_event,
+             async_finish=self.async_finish,
+             allocate_on_comm_stream=(previous_event is not None) and self.async_finish,
+-            expert_alignment=128 if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM else 1,
++            expert_alignment=(
++                128
++                if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
++                and not self.quant_config.get("bf16_weights", False)
++                else 1
++            ),
+             config=DeepEPConfig.get_instance().normal_dispatch_config,
+         )
+         get_global_expert_distribution_recorder().on_deepep_dispatch_normal(
+@@ -491,7 +497,12 @@ class _DeepEPDispatcherImplNormal(_DeepEPDispatcherImplBase):
+         topk_weights: torch.Tensor,
+     ):
+ 
+-        if deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM or _use_aiter or _is_npu:
++        if (
++            deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
++            or _use_aiter
++            or _is_npu
++            or self.quant_config.get("bf16_weights", False)
++        ):
+             output = hidden_states
+         else:
+             raise NotImplementedError()  # triton runner was supported but it's temporarily disabled
+@@ -551,10 +562,12 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
+         buffer = self._get_buffer()
+         topk_weights, topk_ids = topk_output.topk_weights, topk_output.topk_ids
+         topk_ids = topk_ids.to(torch.int64)
+-        expected_m = (
+-            hidden_states.shape[0] * buffer.group_size * topk_ids.shape[1]
+-            + self.num_experts
+-        ) // self.num_experts
++        # Use a correctness-preserving upper bound for per-expert token count.
++        # In the worst case, every rank routes all local tokens to the same expert.
++        expected_m = min(
++            hidden_states.shape[0] * buffer.group_size,
++            self.num_max_dispatch_tokens_per_rank * buffer.group_size,
++        )
+         hidden_states, masked_m, event, hook = self._dispatch_core(
+             hidden_states,
+             topk_ids,
+@@ -609,7 +622,9 @@ class _DeepEPDispatcherImplLowLatency(_DeepEPDispatcherImplBase):
+         input_global_scale = self.quant_config.get("input_global_scale", None)
+         if input_global_scale is not None:
+             use_nvfp4 = True
+-        elif not envs.SGLANG_DEEPEP_BF16_DISPATCH.get():
++        elif not envs.SGLANG_DEEPEP_BF16_DISPATCH.get() and not self.quant_config.get(
++            "bf16_weights", False
++        ):
+             use_fp8 = True
+ 
+         buffer = self._get_buffer()
 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 a13c53af4d..1d80d06b13 100644
+index 4cbfed6f9..88b452744 100644
 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
 +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
-@@ -500,7 +500,7 @@ class CompressedTensorsConfig(QuantizationConfig):
+@@ -499,7 +499,7 @@ class CompressedTensorsConfig(QuantizationConfig):
          )
          is_static = not weight_quant.dynamic
  
@@ -983,7 +1378,7 @@ index a13c53af4d..1d80d06b13 100644
  
      def _is_mxint4a16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool:
          input_quant_none = input_quant is None
-@@ -969,6 +969,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase):
+@@ -968,6 +968,9 @@ class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase):
      def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
          layer.scheme.process_weights_after_loading(layer)
  
@@ -994,7 +1389,7 @@ index a13c53af4d..1d80d06b13 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 7a8fb65421..f1c85899cd 100644
+index 6264f36d0..f0310e305 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 (
@@ -1060,18 +1455,25 @@ index 7a8fb65421..f1c85899cd 100644
          w13_g_idx = torch.nn.Parameter(
              torch.empty(
                  num_experts,
-@@ -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)
+@@ -225,11 +254,14 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
  
+         # 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:
  
-         # Skip if the layer is already converted to Marlin format to prevent double-packing.
-@@ -334,6 +367,24 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
+@@ -334,6 +366,24 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
          )
          replace_tensor("w2_weight_scale", marlin_w2_scales)
  
@@ -1096,7 +1498,7 @@ index 7a8fb65421..f1c85899cd 100644
          layer.is_marlin_converted = True
  
      def restore_weights_before_loading(self, layer: torch.nn.Module):
-@@ -399,6 +450,8 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
+@@ -399,6 +449,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,
@@ -1105,17 +1507,64 @@ index 7a8fb65421..f1c85899cd 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/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py
+index ae0614635..32171c9c1 100644
+--- a/python/sglang/srt/layers/rotary_embedding.py
++++ b/python/sglang/srt/layers/rotary_embedding.py
+@@ -150,9 +150,7 @@ class RotaryEmbedding(MultiPlatformOp):
+ 
+         if get_global_server_args().rl_on_policy_target is not None:
+             self._forward_method = self.forward_native
+-            self._apply_rotary_emb_wrapped = torch.compile(dynamic=True)(
+-                self._apply_rotary_emb_wrapped
+-            )
++
+         self.position_cos, self.position_sin = None, None
+ 
+     def _compute_inv_freq(self, base: Union[int, float]) -> torch.Tensor:
+@@ -1778,6 +1776,9 @@ class MRotaryEmbedding(RotaryEmbedding):
+         key: torch.Tensor,
+         fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
+     ) -> Tuple[torch.Tensor, torch.Tensor]:
++        assert (
++            fused_set_kv_buffer_arg is None
++        ), "fused_set_kv_buffer_arg is not supported for npu implementation"
+         # TODO: remove this when npu_mrope supports QNumHeads * QHeadSize > 4096
+         assert (
+             fused_set_kv_buffer_arg is None
+diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py
+index f78d83d79..f1d577453 100644
+--- a/python/sglang/srt/layers/sampler.py
++++ b/python/sglang/srt/layers/sampler.py
+@@ -122,14 +122,9 @@ class Sampler(nn.Module):
+             # In RL on-policy mode, we use log_softmax to compute logprobs to match the trainer.
+             logprobs_via_logsoftmax_kernel = None
+             if self.rl_on_policy_target is not None:
+-                # TODO: use more inplace ops to save memory
+-                logits_div_temperature = (
+-                    logits.bfloat16().div(sampling_info.temperatures).bfloat16()
+-                )
+                 logprobs_via_logsoftmax_kernel = torch.log_softmax(
+-                    logits_div_temperature, dim=-1
++                    logits / sampling_info.temperatures, dim=-1
+                 )
+-                del logits_div_temperature
+ 
+             if self.use_ascend_backend:
+                 # Ascend backend: sample from logits directly.
 diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py
-index bd97965345..e6a147c1b4 100644
+index ff1774567..42d27a82a 100644
 --- a/python/sglang/srt/managers/io_struct.py
 +++ b/python/sglang/srt/managers/io_struct.py
-@@ -1449,6 +1449,18 @@ class ResumeMemoryOccupationReqOutput(BaseReq):
-     pass
+@@ -1403,6 +1403,20 @@ class UpdateWeightsFromIPCReqOutput(BaseReq):
+     message: str
  
  
 +@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
 +
 +
@@ -1126,62 +1575,24 @@ index bd97965345..e6a147c1b4 100644
 +
 +
  @dataclass
- class CheckWeightsReqInput(BaseReq):
-     action: str
-@@ -1753,6 +1765,8 @@ class GetLoadReqOutput(BaseReq):
+ class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq):
+     success: bool
+@@ -1802,6 +1816,10 @@ 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 e0a1669fb3..fbbb6bb12b 100644
---- a/python/sglang/srt/managers/multi_tokenizer_mixin.py
-+++ b/python/sglang/srt/managers/multi_tokenizer_mixin.py
-@@ -496,6 +496,35 @@ def monkey_patch_uvicorn_multiprocessing(timeout: float = 10):
-             "uvicorn.supervisors.multiprocess not found, skipping monkey patch"
-         )
- 
-+    try:
-+        import uvicorn._subprocess as uvicorn_subprocess
-+        import uvicorn.supervisors.multiprocess as uvicorn_multiprocess
-+
-+        def _safe_get_stdin_fileno():
-+            try:
-+                fileno = sys.stdin.fileno()
-+                dup_fd = os.dup(fileno)
-+                os.close(dup_fd)
-+                return fileno
-+            except (AttributeError, OSError):
-+                return None
-+
-+        def _patched_get_subprocess(config, target, sockets):
-+            kwargs = {
-+                "config": config,
-+                "target": target,
-+                "sockets": sockets,
-+                "stdin_fileno": _safe_get_stdin_fileno(),
-+            }
-+            return uvicorn_subprocess.spawn.Process(
-+                target=uvicorn_subprocess.subprocess_started, kwargs=kwargs
-+            )
-+
-+        uvicorn_subprocess.get_subprocess = _patched_get_subprocess
-+        uvicorn_multiprocess.get_subprocess = _patched_get_subprocess
-+    except Exception:
-+        pass
-+
- 
- 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 0b26be6c6d..2ea1042cf9 100644
+index c07995798..dd8ca7167 100644
 --- a/python/sglang/srt/managers/schedule_batch.py
 +++ b/python/sglang/srt/managers/schedule_batch.py
-@@ -1972,7 +1972,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
+@@ -1869,7 +1869,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
          while first_iter or (
              not self.check_decode_mem(selected_indices=sorted_indices)
          ):
@@ -1194,18 +1605,18 @@ index 0b26be6c6d..2ea1042cf9 100644
                  break
  
 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
-index 67af2d0de9..122ddb3874 100644
+index a9ff0ac94..ba0a75ee7 100644
 --- a/python/sglang/srt/managers/scheduler.py
 +++ b/python/sglang/srt/managers/scheduler.py
-@@ -120,6 +120,7 @@ from sglang.srt.managers.io_struct import (
-     LoadLoRAAdapterReqOutput,
+@@ -114,6 +114,7 @@ from sglang.srt.managers.io_struct import (
      OpenSessionReqInput,
+     OpenSessionReqOutput,
      PauseGenerationReqInput,
 +    PostProcessWeightsReqInput,
      ProfileReq,
      ReleaseMemoryOccupationReqInput,
      ResumeMemoryOccupationReqInput,
-@@ -1232,6 +1233,7 @@ class Scheduler(
+@@ -1063,6 +1064,7 @@ class Scheduler(
                  ),
                  (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
                  (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc),
@@ -1213,34 +1624,284 @@ index 67af2d0de9..122ddb3874 100644
                  (GetWeightsByNameReqInput, self.get_weights_by_name),
                  (ReleaseMemoryOccupationReqInput, self.release_memory_occupation),
                  (ResumeMemoryOccupationReqInput, self.resume_memory_occupation),
+@@ -1304,7 +1306,6 @@ class Scheduler(
+                 self.tp_cpu_group,
+                 src=self.tp_group.ranks[0],
+             )
+-
+         # Process MM requests under EPD-disaggregation mode
+         if (
+             self.pp_rank == 0
+diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py
+index 30b2732b9..68090b161 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 496cd96656..cf2d43015a 100644
+index 482bc6ca6..857cfa6a3 100644
 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py
 +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py
-@@ -1154,7 +1154,7 @@ class SchedulerOutputProcessorMixin:
-         dp_ranks = [self.dp_rank] * len(rids) if rids else None
+@@ -1134,7 +1134,7 @@ class SchedulerOutputProcessorMixin:
+                 req.log_time_stats()
  
          # 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(
-                 BatchTokenIDOutput(
-                     rids=rids,
-diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py
-index c02ed7997d..61733c4127 100644
---- a/python/sglang/srt/managers/scheduler_profiler_mixin.py
-+++ b/python/sglang/srt/managers/scheduler_profiler_mixin.py
-@@ -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)
--            elif batch.forward_mode.is_decode():
-+            elif batch.forward_mode.is_decode() or batch.forward_mode.is_prebuilt():
-                 if self.profiler_decode_ct == 0:
-                     if self.profile_in_progress:
-                         # force trace flush
+diff --git a/python/sglang/srt/managers/scheduler_pp_mixin.py b/python/sglang/srt/managers/scheduler_pp_mixin.py
+index 1a65a3c3d..245722295 100644
+--- a/python/sglang/srt/managers/scheduler_pp_mixin.py
++++ b/python/sglang/srt/managers/scheduler_pp_mixin.py
+@@ -20,6 +20,7 @@ from sglang.srt.layers.dp_attention import (
+     get_attention_dp_rank,
+     get_attention_dp_size,
+     is_dp_attention_enabled,
++    set_is_extend_in_batch,
+ )
+ from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+ from sglang.srt.managers.utils import (
+@@ -224,7 +225,28 @@ class SchedulerPPMixin:
+ 
+                 self.process_prefill_chunk()
+                 batch = self.get_new_batch_prefill()
+-                batch = self.maybe_prepare_mlp_sync_batch(batch)
++                need_mlp_sync = self.require_mlp_sync
++                skipped_mlp_sync = False
++                if (
++                    need_mlp_sync
++                    and self.disaggregation_mode == DisaggregationMode.PREFILL
++                    and self.server_args.enable_nsa_prefill_context_parallel
++                    and self.pp_size > 1
++                ):
++                    # In PD prefill CP+PP, MLP sync all_gather can deadlock on idle micro-batches.
++                    # Skip MLP sync here because decode-side MLP gather is not involved in this path.
++                    need_mlp_sync = False
++                    skipped_mlp_sync = True
++                batch = self.maybe_prepare_mlp_sync_batch(
++                    batch, need_sync=need_mlp_sync
++                )
++                if skipped_mlp_sync:
++                    # MLP sync was skipped but set_is_extend_in_batch is still needed
++                    # by the deepep dispatcher (called in model forward).
++                    is_extend = (
++                        batch.forward_mode.is_extend() if batch is not None else False
++                    )
++                    set_is_extend_in_batch(is_extend)
+                 self.mbs[mb_id] = batch
+                 self.running_mbs[mb_id] = self.running_batch
+ 
+@@ -288,6 +310,11 @@ class SchedulerPPMixin:
+                         next_batch_result,
+                     )
+                     self.last_mbs[next_mb_id] = self.mbs[next_mb_id]
++                    if self.current_scheduler_metrics_enabled:
++                        self.log_prefill_stats(
++                            prefill_stats=self.mbs[next_mb_id].prefill_stats,
++                            can_run_cuda_graph=next_batch_result.can_run_cuda_graph,
++                        )
+ 
+                 if tmbs[next_mb_id] is not None:
+                     self.process_disagg_prefill_inflight_queue(next_release_rids)
+@@ -524,6 +551,11 @@ class SchedulerPPMixin:
+         self.last_rank_comm_queue: deque[Tuple[torch.cuda.Event, PPProxyTensors]] = (
+             deque()
+         )
++        # PP1 (last rank) stores its own batch outputs locally to avoid the
++        # PP1→PP0→PP1 round-trip that causes a deadlock in disagg prefill.
++        self.last_rank_local_result_queue: deque[Tuple[torch.cuda.Event, PPProxyTensors]] = (
++            deque()
++        )
+ 
+         self.send_req_work = []
+         self.send_proxy_work = []
+@@ -859,31 +891,39 @@ class SchedulerPPMixin:
+ 
+     def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False):
+         p2p_work = []
+-        if self.attn_tp_rank == 0:
+-            dp_offset = self.attn_dp_rank * self.attn_tp_size
++        if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
++            lane_offset = self.attn_dp_rank * self.attn_tp_size
+             p2p_work = point_to_point_pyobj(
+                 data,
+-                self.pp_rank * self.tp_size + dp_offset,
++                self.pp_rank * self.tp_size + lane_offset,
+                 self.world_group.cpu_group,
+-                self.pp_rank * self.tp_size + dp_offset,
+-                ((self.pp_rank + 1) % self.pp_size) * self.tp_size + dp_offset,
++                self.pp_rank * self.tp_size + lane_offset,
++                ((self.pp_rank + 1) % self.pp_size) * self.tp_size + lane_offset,
+                 async_send=async_send,
+             )
+         return p2p_work
+ 
+     def _pp_recv_pyobj_from_prev_stage(self: Scheduler):
+-        if self.attn_tp_rank == 0:
+-            dp_offset = self.attn_dp_rank * self.attn_tp_size
++        if self.attn_tp_rank == 0 and self.attn_cp_rank == 0:
++            lane_offset = self.attn_dp_rank * self.attn_tp_size
+             data = point_to_point_pyobj(
+                 [],
+-                self.pp_rank * self.tp_size + dp_offset,
++                self.pp_rank * self.tp_size + lane_offset,
+                 self.world_group.cpu_group,
+-                ((self.pp_rank - 1) % self.pp_size) * self.tp_size + dp_offset,
+-                self.pp_rank * self.tp_size + dp_offset,
++                ((self.pp_rank - 1) % self.pp_size) * self.tp_size + lane_offset,
++                self.pp_rank * self.tp_size + lane_offset,
+             )
+         else:
+             data = None
+ 
++        if self.attn_cp_size > 1:
++            data = broadcast_pyobj(
++                data,
++                self.attn_cp_group.rank,
++                self.attn_cp_cpu_group,
++                src=self.attn_cp_group.ranks[0],
++            )
++
+         if self.attn_tp_size > 1:
+             data = broadcast_pyobj(
+                 data,
+@@ -1004,8 +1044,13 @@ class SchedulerPPMixin:
+                             pp_outputs_to_send.tensors,
+                             async_send=True,
+                         )
+-        # send the outputs from the last round to let the next stage worker run post processing
+-        if not self.pp_group.is_last_rank:
++                # Store locally so the last rank can process its own batch result
++                # without receiving from the second-to-last rank (avoids deadlock).
++                self.last_rank_local_result_queue.append((q_event, pp_outputs_to_send))
++        elif self.pp_rank != self.pp_size - 2:
++            # Forward output through the chain: PP0→PP1→...→PP(last-2).
++            # The second-to-last rank does NOT forward to the last rank because
++            # the last rank uses last_rank_local_result_queue instead of receiving.
+             if pp_outputs:
+                 with torch.profiler.record_function("send_res_dict_to_next_stage"):
+                     send_output_work = self._pp_send_dict_to_next_stage(
+@@ -1034,20 +1079,36 @@ class SchedulerPPMixin:
+         )
+ 
+         if mbs[next_mb_id] is not None:
+-            with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
+-                next_pp_outputs = None
++            if self.pp_group.is_last_rank:
++                # Last rank: use the locally-stored output instead of receiving
++                # from the second-to-last rank. Receiving would cause a deadlock
++                # because the chain PP_last→PP0→...→PP(last-2)→PP_last requires
++                # PP0 to have pp_outputs ready, which it doesn't on the first batch.
+                 if not mbs[next_mb_id].forward_mode.is_prebuilt():
+-                    next_pp_outputs = PPProxyTensors(
+-                        self._pp_recv_dict_from_prev_stage()
+-                    )
+-            if not mbs[next_mb_id].forward_mode.is_prebuilt():
+-                with self.copy_stream_ctx:
+-                    self.copy_stream.wait_stream(self.default_stream)
+-                    batch_result = self._pp_prep_batch_result(
+-                        mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
+-                    )
+-                    d2h_event = torch.cuda.Event()
+-                    d2h_event.record(torch.cuda.current_stream())
++                    q_event, next_pp_outputs = self.last_rank_local_result_queue.popleft()
++                    with self.copy_stream_ctx:
++                        torch.cuda.current_stream().wait_event(q_event)
++                        self.copy_stream.wait_stream(self.default_stream)
++                        batch_result = self._pp_prep_batch_result(
++                            mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
++                        )
++                        d2h_event = torch.cuda.Event()
++                        d2h_event.record(torch.cuda.current_stream())
++            else:
++                with torch.profiler.record_function("recv_res_dict_from_prev_stage"):
++                    next_pp_outputs = None
++                    if not mbs[next_mb_id].forward_mode.is_prebuilt():
++                        next_pp_outputs = PPProxyTensors(
++                            self._pp_recv_dict_from_prev_stage()
++                        )
++                if not mbs[next_mb_id].forward_mode.is_prebuilt():
++                    with self.copy_stream_ctx:
++                        self.copy_stream.wait_stream(self.default_stream)
++                        batch_result = self._pp_prep_batch_result(
++                            mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
++                        )
++                        d2h_event = torch.cuda.Event()
++                        d2h_event.record(torch.cuda.current_stream())
+ 
+         return next_pp_outputs, batch_result, d2h_event, send_output_work
+ 
+@@ -1085,9 +1146,12 @@ class SchedulerPPMixin:
+         """
+         Used by PP, get the required rids with the given poll statuses.
+         """
++        gloo_group = self.attn_tp_cpu_group
++        if self.attn_cp_size > 1:
++            gloo_group = self.tp_cpu_group
+         polls = poll_and_all_reduce(
+             [req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue],
+-            self.attn_tp_cpu_group,
++            gloo_group,
+         )
+         rids: List = []
+         for poll_statuses in poll_statuses_group:
 diff --git a/python/sglang/srt/managers/scheduler_update_weights_mixin.py b/python/sglang/srt/managers/scheduler_update_weights_mixin.py
-index abcda67946..a53848b79d 100644
+index 293a84350..244ea4eb1 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 (
@@ -1260,7 +1921,7 @@ index abcda67946..a53848b79d 100644
      ReleaseMemoryOccupationReqInput,
      ReleaseMemoryOccupationReqOutput,
      ResumeMemoryOccupationReqInput,
-@@ -117,6 +120,11 @@ class SchedulerUpdateWeightsMixin:
+@@ -114,6 +117,11 @@ class SchedulerUpdateWeightsMixin:
          torch.distributed.barrier(group=self.tp_cpu_group)
          return UpdateWeightsFromIPCReqOutput(success, message)
  
@@ -1272,7 +1933,7 @@ index abcda67946..a53848b79d 100644
      def get_weights_by_name(self: Scheduler, recv_req: GetWeightsByNameReqInput):
          parameter = self.tp_worker.get_weights_by_name(recv_req)
          return GetWeightsByNameReqOutput(parameter)
-@@ -140,6 +148,15 @@ class SchedulerUpdateWeightsMixin:
+@@ -137,6 +145,15 @@ class SchedulerUpdateWeightsMixin:
              self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
              self.flush_cache()
  
@@ -1288,7 +1949,7 @@ index abcda67946..a53848b79d 100644
          if GPU_MEMORY_TYPE_WEIGHTS in tags:
              self.stashed_model_static_state = _export_static_state(
                  self.tp_worker.model_runner.model
-@@ -180,6 +197,15 @@ class SchedulerUpdateWeightsMixin:
+@@ -177,6 +194,15 @@ class SchedulerUpdateWeightsMixin:
          if GPU_MEMORY_TYPE_KV_CACHE in tags:
              self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
  
@@ -1305,7 +1966,7 @@ index abcda67946..a53848b79d 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 544c609401..841658c30e 100644
+index f2ffa9909..6e4d1d460 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 (
@@ -1338,7 +1999,7 @@ index 544c609401..841658c30e 100644
                  (
                      GetWeightsByNameReqOutput,
                      self.get_weights_by_name_communicator.handle_recv,
-@@ -530,6 +539,17 @@ class TokenizerCommunicatorMixin:
+@@ -522,6 +531,17 @@ class TokenizerCommunicatorMixin:
  
          return success, message
  
@@ -1354,13 +2015,27 @@ index 544c609401..841658c30e 100644
 +            return _Communicator.merge_results(results)
 +
      async def init_weights_send_group_for_remote_instance(
-         self: TokenizerManager,
+         self,
          obj: InitWeightsSendGroupForRemoteInstanceReqInput,
 diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
-index 81424329a0..2c132be63d 100644
+index 0914a5230..cce2d8a2b 100644
 --- a/python/sglang/srt/managers/tokenizer_manager.py
 +++ b/python/sglang/srt/managers/tokenizer_manager.py
-@@ -1383,7 +1383,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin):
+@@ -324,8 +324,12 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
+             context, zmq.PULL, port_args.tokenizer_ipc_name, True
+         )
+         if self.server_args.tokenizer_worker_num == 1:
++            self.send_to_scheduler_context = zmq.Context(1)
+             self.send_to_scheduler = get_zmq_socket(
+-                context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
++                self.send_to_scheduler_context,
++                zmq.PUSH,
++                port_args.scheduler_input_ipc_name,
++                True,
+             )
+         else:
+             from sglang.srt.managers.multi_tokenizer_mixin import SenderWrapper
+@@ -1327,7 +1331,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
          async with self.is_pause_cond:
              self.is_pause = True
              if obj.mode != "abort":
@@ -1369,7 +2044,7 @@ index 81424329a0..2c132be63d 100644
              else:
                  # we are using the model_update_lock to check if there is still on-going requests.
                  while True:
-@@ -1397,7 +1397,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin):
+@@ -1341,7 +1345,7 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerMultiItemMixi
      async def continue_generation(self, obj: ContinueGenerationReqInput):
          async with self.is_pause_cond:
              self.is_pause = False
@@ -1378,41 +2053,8 @@ index 81424329a0..2c132be63d 100644
              self.is_pause_cond.notify_all()
  
      async def update_weights_from_disk(
-@@ -1965,25 +1965,23 @@ class TokenizerManager(TokenizerCommunicatorMixin, TokenizerManagerScoreMixin):
-             priority = getattr(state.obj, "priority", None)
-             if priority is not None:
-                 labels["priority"] = str(priority)
--        if (
--            not state.ttft_observed
--            and self.disaggregation_mode != DisaggregationMode.PREFILL
--        ):
-+        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.time_stats.get_first_token_latency()
--            )
-+            if self.disaggregation_mode != DisaggregationMode.PREFILL:
-+                self.metrics_collector.observe_time_to_first_token(
-+                    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:
-                 self.metrics_collector.observe_inter_token_latency(
-                     labels,
-                     state.time_stats.get_interval(),
-                     num_new_tokens,
-                 )
-                 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 7f63610da8..fb56de1583 100644
+index 86b009df4..16ebd52ae 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 (
@@ -1423,7 +2065,7 @@ index 7f63610da8..fb56de1583 100644
      SendWeightsToRemoteInstanceReqInput,
      UnloadLoRAAdapterReqInput,
      UpdateWeightFromDiskReqInput,
-@@ -170,6 +171,11 @@ class BaseTpWorker(ABC):
+@@ -168,6 +169,11 @@ class BaseTpWorker(ABC):
          success, message = self.model_runner.update_weights_from_ipc(recv_req)
          return success, message
  
@@ -1435,11 +2077,135 @@ index 7f63610da8..fb56de1583 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 fa08bb66a..fa539315c 100644
+--- a/python/sglang/srt/mem_cache/allocator.py
++++ b/python/sglang/srt/mem_cache/allocator.py
+@@ -347,6 +347,84 @@ def alloc_decode_kernel(
+         tl.store(out_indices + pid, page * page_size)
+ 
+ 
++def alloc_extend_torch_fallback(
++    prefix_lens_cpu: torch.Tensor,
++    seq_lens_cpu: torch.Tensor,
++    last_loc: torch.Tensor,
++    free_pages: torch.Tensor,
++    out_indices: torch.Tensor,
++    page_size: int,
++    debug_mode: bool = False,
++):
++    extend_lens_cpu = (seq_lens_cpu - prefix_lens_cpu).to(torch.int64)
++    if extend_lens_cpu.numel() == 0:
++        return
++
++    output_start_locs_cpu = torch.cumsum(extend_lens_cpu, dim=0) - extend_lens_cpu
++    num_pages_after = (seq_lens_cpu + page_size - 1) // page_size
++    num_pages_before = (prefix_lens_cpu + page_size - 1) // page_size
++    num_new_pages_cpu = num_pages_after - num_pages_before
++    page_start_locs_cpu = torch.cumsum(num_new_pages_cpu, dim=0) - num_new_pages_cpu
++
++    total_new_pages = int(num_new_pages_cpu.sum().item())
++    if total_new_pages > free_pages.numel():
++        return
++
++    if debug_mode:
++        assert int(extend_lens_cpu.sum().item()) == out_indices.numel()
++
++    prefix_lens_list = prefix_lens_cpu.tolist()
++    seq_lens_list = seq_lens_cpu.tolist()
++    extend_lens_list = extend_lens_cpu.tolist()
++    out_start_list = output_start_locs_cpu.tolist()
++    page_start_list = page_start_locs_cpu.tolist()
++    num_new_pages_list = num_new_pages_cpu.tolist()
++
++    device = out_indices.device
++    dtype = out_indices.dtype
++    offsets_page = torch.arange(page_size, device=device, dtype=dtype)
++
++    for i, extend_len in enumerate(extend_lens_list):
++        if extend_len == 0:
++            continue
++
++        pre_len = prefix_lens_list[i]
++        seq_len = seq_lens_list[i]
++        out_start = out_start_list[i]
++        page_start = page_start_list[i]
++        num_new_pages = num_new_pages_list[i]
++
++        pre_mod = pre_len % page_size
++        part1 = min(extend_len, page_size - pre_mod) if pre_mod != 0 else 0
++        if part1:
++            start_val = last_loc[i] + 1
++            out_indices[out_start : out_start + part1] = start_val + torch.arange(
++                part1, device=device, dtype=dtype
++            )
++            if part1 == extend_len:
++                continue
++
++        ceil_pre_pages = (pre_len + page_size - 1) // page_size
++        full_pages_after = seq_len // page_size
++        num_full_pages = full_pages_after - ceil_pre_pages
++        if num_full_pages < 0:
++            num_full_pages = 0
++        part2 = num_full_pages * page_size
++        if part2:
++            pages = free_pages[page_start : page_start + num_full_pages]
++            full_indices = (pages[:, None] * page_size + offsets_page).reshape(-1)
++            out_indices[out_start + part1 : out_start + part1 + part2] = full_indices
++            if part1 + part2 == extend_len:
++                continue
++
++        part3 = extend_len - part1 - part2
++        if part3:
++            last_page = free_pages[page_start + num_new_pages - 1]
++            out_indices[out_start + part1 + part2 : out_start + extend_len] = (
++                last_page * page_size + torch.arange(part3, device=device, dtype=dtype)
++            )
++
++
+ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
+     """
+     An allocator managing the indices to kv cache data.
+@@ -411,7 +489,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 +502,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 3c1e97daab..e5128e5ee2 100644
+index d7cd472a9..1207ab199 100644
 --- a/python/sglang/srt/mem_cache/hiradix_cache.py
 +++ b/python/sglang/srt/mem_cache/hiradix_cache.py
-@@ -755,9 +755,8 @@ class HiRadixCache(RadixCache):
+@@ -76,6 +76,7 @@ class HiRadixCache(RadixCache):
+                 allocator_type=server_args.hicache_storage_backend,
+             )
+         elif isinstance(self.kv_cache, NSATokenToKVPool):
++            # Check NSA before MLA since NSATokenToKVPool is a subclass of MLATokenToKVPool
+             self.token_to_kv_pool_host = NSATokenToKVPoolHost(
+                 self.kv_cache,
+                 server_args.hicache_ratio,
+@@ -94,7 +95,7 @@ class HiRadixCache(RadixCache):
+                 allocator_type=server_args.hicache_storage_backend,
+             )
+         else:
+-            raise ValueError(f"HiRadixCache only supports MHA and MLA yet")
++            raise ValueError(f"HiRadixCache only supports MHA and MLA and NSA yet")
+ 
+         self.tp_group = params.tp_cache_group
+         self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
+@@ -750,9 +751,8 @@ class HiRadixCache(RadixCache):
              self._update_leaf_status(node)
              self._update_host_leaf_status(node)
              if node.parent is None:
@@ -1449,37 +2215,13 @@ index 3c1e97daab..e5128e5ee2 100644
 +                # Node belongs to a stale (flushed) tree — stop traversal gracefully.
 +                break
              node = node.parent
-         return DecLockRefResult(delta=delta)
- 
-@@ -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)
-+        self._update_host_leaf_status(node.parent)
-         return num_evicted
- 
-     def _evict_regular(self, node: TreeNode):
-@@ -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)
-+                    self._update_host_leaf_status(node.parent)
-                 else:
-                     self._inc_hit_count(node, chunked)
-                     total_prefix_length += prefix_len
-@@ -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)
-+                    self._update_host_leaf_status(new_node.parent)
-                 else:
-                     self._inc_hit_count(new_node, chunked)
-                     total_prefix_length += prefix_len
+         return delta
+ 
 diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py
-index e4c158cda9..cf7333235f 100644
+index 1d917137c..669e5c518 100644
 --- a/python/sglang/srt/mem_cache/memory_pool.py
 +++ b/python/sglang/srt/mem_cache/memory_pool.py
-@@ -1854,9 +1854,12 @@ class NSATokenToKVPool(MLATokenToKVPool):
+@@ -1777,9 +1777,12 @@ class NSATokenToKVPool(MLATokenToKVPool):
          else:
              assert self.page_size == 64
          with (
@@ -1495,19 +2237,19 @@ index e4c158cda9..cf7333235f 100644
          ):
              self.index_k_with_scale_buffer = [
                  torch.zeros(
-@@ -1878,6 +1881,11 @@ class NSATokenToKVPool(MLATokenToKVPool):
+@@ -1801,6 +1804,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:
-@@ -1960,6 +1968,50 @@ class NSATokenToKVPool(MLATokenToKVPool):
+@@ -1876,6 +1884,50 @@ class NSATokenToKVPool(MLATokenToKVPool):
          ]
          return data_ptrs, data_lens, item_lens
  
@@ -1559,10 +2301,10 @@ index e4c158cda9..cf7333235f 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 7d16160372..70fbdc702f 100644
+index 42b169728..fbefb0193 100644
 --- a/python/sglang/srt/mem_cache/radix_cache.py
 +++ b/python/sglang/srt/mem_cache/radix_cache.py
-@@ -512,7 +512,17 @@ class RadixCache(BasePrefixCache):
+@@ -495,7 +495,13 @@ class RadixCache(BasePrefixCache):
          if self.disable:
              return
  
@@ -1573,15 +2315,11 @@ index 7d16160372..70fbdc702f 100644
 +        # req_to_token_pool initialization), leading to spurious tree nodes and memory
 +        # leak when page-aligned token counts happen to cross a page boundary.
 +        kv_committed_len = req.kv_committed_len
-+        token_ids = (
-+            req.fill_ids[:kv_committed_len]
-+            if kv_committed_len < len(req.fill_ids)
-+            else req.fill_ids
-+        )
++        token_ids = req.fill_ids[:kv_committed_len] if kv_committed_len < len(req.fill_ids) else req.fill_ids
          kv_indices = self.req_to_token_pool.req_to_token[
              req.req_pool_idx, : len(token_ids)
          ]
-@@ -638,9 +648,8 @@ class RadixCache(BasePrefixCache):
+@@ -619,9 +625,8 @@ class RadixCache(BasePrefixCache):
              node.lock_ref -= 1
              self._update_leaf_status(node)
              if node.parent is None:
@@ -1591,13 +2329,13 @@ index 7d16160372..70fbdc702f 100644
 +                # Node belongs to a stale (flushed) tree — stop traversal gracefully.
 +                break
              node = node.parent
-         return DecLockRefResult(delta=delta)
+         return delta
  
 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
-index a59742b943..a7347c15b8 100644
+index 275775a73..c67f3342e 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):
+@@ -395,7 +395,10 @@ class ModelRunner(ModelRunnerKVCacheMixin):
          self.forward_stream = torch.get_device_module(self.device).Stream()
  
          # CPU offload
@@ -1605,14 +2343,12 @@ index a59742b943..a7347c15b8 100644
 +        # 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)
-+            )
++            set_offloader(create_offloader_from_server_args(server_args, dp_rank=dp_rank))
  
          self._weight_checker = WeightChecker(model_runner=self)
  
-@@ -646,7 +651,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
-             )
+@@ -600,7 +603,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
+         )
  
          # Init routed experts capturer
 -        self.init_routed_experts_capturer()
@@ -1621,7 +2357,7 @@ index a59742b943..a7347c15b8 100644
  
          if self.device == "cuda" or self.device == "musa":
              self.init_cublas()
-@@ -2767,11 +2773,19 @@ class ModelRunner(ModelRunnerKVCacheMixin):
+@@ -2429,11 +2433,19 @@ class ModelRunner(ModelRunnerKVCacheMixin):
          output.expert_distribution_metrics = recorder_outputs.get("metrics")
  
          # Copy cached routing experts' buffers back to CPU cache
@@ -1646,7 +2382,7 @@ index a59742b943..a7347c15b8 100644
  
          if self.eplb_manager is not None:
              self.eplb_manager.on_forward_pass_end()
-@@ -3021,6 +3035,42 @@ class ModelRunner(ModelRunnerKVCacheMixin):
+@@ -2664,6 +2676,42 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                      device=self.device,
                  )
  
@@ -1689,129 +2425,341 @@ index a59742b943..a7347c15b8 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/glm4v_moe.py b/python/sglang/srt/models/glm4v_moe.py
-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):
-         self.num_fused_shared_experts = 0
-         self.determine_num_fused_shared_experts()
- 
--        self.model = Glm4MoeModel(
--            config,
--            quant_config,
--            prefix=add_prefix("language_model", prefix),
--        )
-+        if not self.config.encoder_only:
-+            self.model = Glm4MoeModel(
-+                config,
-+                quant_config,
-+                prefix=add_prefix("language_model", prefix),
-+            )
+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 cc673a9ca..06c430d2c 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/gpt_oss.py b/python/sglang/srt/models/gpt_oss.py
+index 2cf813bce..1250c49e4 100644
+--- a/python/sglang/srt/models/gpt_oss.py
++++ b/python/sglang/srt/models/gpt_oss.py
+@@ -17,6 +17,7 @@
+ 
+ import logging
+ import math
++import re
+ from collections.abc import Iterable
+ from functools import partial
+ from typing import Any, Dict, List, Optional, Tuple, Union
+@@ -1065,6 +1066,12 @@ class GptOssForCausalLM(nn.Module):
+                 weight_loader(param, loaded_weight, shard_id)
+                 break
+             else:
++                # Try per-expert format: experts.{id}.{gate_proj|up_proj|down_proj}.{weight|bias}
++                per_expert_match = _PER_EXPERT_RE.match(name)
++                if per_expert_match:
++                    _load_per_expert_param(per_expert_match, loaded_weight, params_dict)
++                    continue
 +
-+            if self.pp_group.is_last_rank:
-+                if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
-+                    self.lm_head = self.model.embed_tokens
-+                else:
-+                    self.lm_head = ParallelLMHead(
-+                        config.vocab_size,
-+                        config.hidden_size,
-+                        quant_config=quant_config,
-+                        prefix=add_prefix("lm_head", prefix),
-+                        use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
-+                    )
-+            else:
-+                # ranks other than the last rank will have a placeholder layer
-+                self.lm_head = PPMissingLayer()
-+        else:
-+            # encoder_only mode: no language model, so no lm_head needed
-+            self.lm_head = None
-+
-         self.visual = Glm4vVisionModel(
-             config.vision_config,
-             quant_config=quant_config,
-@@ -64,24 +84,14 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-             use_data_parallel=self.use_data_parallel,
-         )
+                 for mapping in expert_params_mapping:
+                     param_name, weight_name, shard_id = mapping
+                     if weight_name not in name:
+@@ -1143,6 +1150,88 @@ class GptOssForCausalLM(nn.Module):
+         return get_attention_sliding_window_size(self.config)
  
--        if self.pp_group.is_last_rank:
--            if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
--                self.lm_head = self.model.embed_tokens
--            else:
--                self.lm_head = ParallelLMHead(
--                    config.vocab_size,
--                    config.hidden_size,
--                    quant_config=quant_config,
--                    prefix=add_prefix("lm_head", prefix),
--                    use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
--                )
--        else:
--            # ranks other than the last rank will have a placeholder layer
--            self.lm_head = PPMissingLayer()
--
-         self.logits_processor = LogitsProcessor(config)
-         self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
--        self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling
-+        _rope_cfg = (
-+            getattr(self.config, "rope_scaling", None)
-+            or getattr(self.config, "rope_parameters", None)
-+            or {}
-+        )
-+        self.is_mrope_enabled = "mrope_section" in _rope_cfg
  
-         # For EAGLE3 support
-         self.capture_aux_hidden_states = False
-@@ -219,6 +229,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                 # Skip loading extra bias for GPTQ models.
-                 if name.endswith(".bias") and name not in params_dict:
-                     continue
-+                # Skip loading visual/language model weights
-+                if (
-+                    self.config.encoder_only or self.config.language_only
-+                ) and name not in params_dict:
-+                    continue
-                 if name not in params_dict:
-                     continue
++# Regex for per-expert weight names: model.layers.X.mlp.experts.E.{proj}.{weight|bias}
++_PER_EXPERT_RE = re.compile(
++    r"(.+\.mlp\.experts\.)(\d+)\.(gate_proj|up_proj|down_proj)\.(weight|bias)"
++)
++
++
++def _load_per_expert_param(match, loaded_weight, params_dict):
++    """Load a per-expert weight/bias tensor into the fused FusedMoE parameter.
++
++    Handles the mapping from per-expert names (e.g., experts.0.gate_proj.weight)
++    to fused parameters (e.g., experts.w13_weight).
++    """
++    prefix, eid_str, proj, ptype = match.groups()
++    eid = int(eid_str)
++
++    # Determine target fused parameter name
++    if proj in ("gate_proj", "up_proj"):
++        key = prefix + ("w13_weight" if ptype == "weight" else "w13_weight_bias")
++    else:  # down_proj
++        key = prefix + ("w2_weight" if ptype == "weight" else "w2_weight_bias")
++
++    if key not in params_dict:
++        return
++
++    param = params_dict[key]
++    expert_slice = param.data[eid]  # slice for this expert
++
++    # Detect triton transposed layout from shape:
++    # w13: transposed=(E, hidden, 2*inter), non-transposed=(E, 2*inter, hidden)
++    # For w13, the larger dim is 2*intermediate; if it's dim -1, layout is transposed.
++    is_transposed = getattr(param, "is_transposed", False)
++    if not is_transposed and ptype == "weight" and "w13" in key:
++        # Infer from shape: transposed has shape[-1] > shape[-2]
++        is_transposed = param.data.shape[-1] > param.data.shape[-2]
++
++    if ptype == "weight":
++        if proj in ("gate_proj", "up_proj"):
++            # w13_weight: gate in first half, up in second half
++            if is_transposed:
++                # Triton layout: (E, hidden, 2*intermediate)
++                half = expert_slice.shape[1] // 2
++                dst = (
++                    expert_slice[:, :half]
++                    if proj == "gate_proj"
++                    else expert_slice[:, half:]
++                )
++            else:
++                # Standard layout: (E, 2*intermediate, hidden)
++                half = expert_slice.shape[0] // 2
++                dst = (
++                    expert_slice[:half] if proj == "gate_proj" else expert_slice[half:]
++                )
++            # loaded_weight shape: (intermediate, hidden)
++            if is_transposed:
++                dst.copy_(loaded_weight.t())
++            else:
++                dst.copy_(loaded_weight)
++        else:
++            # w2_weight: loaded_weight shape (hidden, intermediate)
++            # Detect transposition for w2 as well
++            w2_transposed = is_transposed
++            if not w2_transposed:
++                w2_transposed = param.data.shape[-1] > param.data.shape[-2]
++            if w2_transposed:
++                expert_slice.copy_(loaded_weight.t())
++            else:
++                expert_slice.copy_(loaded_weight)
++    else:
++        # Bias handling
++        if proj in ("gate_proj", "up_proj"):
++            # w13_weight_bias: (E, 2*intermediate)
++            half = expert_slice.shape[0] // 2
++            dst = expert_slice[:half] if proj == "gate_proj" else expert_slice[half:]
++            dst.copy_(loaded_weight)
++        else:
++            # w2_weight_bias: (E, hidden) - only rank 0 loads, others zero
++            if get_moe_tensor_parallel_rank() == 0:
++                expert_slice.copy_(loaded_weight)
++            else:
++                expert_slice.zero_()
++
++
+ def _canonicalize_weights(config, weights_in: Iterable[Tuple[str, torch.Tensor]]):
+     weights_out_dict = dict(weights_in)
  
-@@ -234,6 +249,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                     param_name, weight_name, expert_id, shard_id = mapping
-                     if weight_name not in name:
-                         continue
-+                    if "visual" in name or self.config.encoder_only:
-+                        continue
- 
-                     # Mark as expert weight regardless of whether we can process it
-                     is_expert_weight = True
-@@ -265,6 +282,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                     # Skip loading extra bias for GPTQ models.
-                     if name.endswith(".bias") and name not in params_dict:
-                         continue
-+                    # Skip loading mm/language parameters
-+                    if (
-+                        self.config.encoder_only or self.config.language_only
-+                    ) and name not in params_dict:
-+                        continue
-                     if name not in params_dict:
-                         continue
+diff --git a/python/sglang/srt/models/qwen2.py b/python/sglang/srt/models/qwen2.py
+index a3cfde4d6..18fb07b9d 100644
+--- a/python/sglang/srt/models/qwen2.py
++++ b/python/sglang/srt/models/qwen2.py
+@@ -91,9 +91,6 @@ class Qwen2MLP(nn.Module):
+         self.act_fn = SiluAndMul()
  
+     def forward(self, x):
+-        if get_global_server_args().rl_on_policy_target is not None:
+-            x = x.bfloat16()
+-
+         gate_up, _ = self.gate_up_proj(x)
+         x = self.act_fn(gate_up)
+         x, _ = self.down_proj(x)
+@@ -280,11 +277,6 @@ class Qwen2Model(nn.Module):
+                 quant_config=quant_config,
+                 use_attn_tp_group=is_dp_attention_enabled(),
+                 prefix=add_prefix("embed_tokens", prefix),
+-                params_dtype=(
+-                    torch.float32
+-                    if get_global_server_args().rl_on_policy_target is not None
+-                    else None
+-                ),
+             )
+         else:
+             self.embed_tokens = PPMissingLayer()
+@@ -307,10 +299,8 @@ class Qwen2Model(nn.Module):
+         if self.pp_group.is_last_rank:
+             norm_kwargs = (
+                 dict(
+-                    weight_dtype=torch.float32,
+                     cast_x_before_out_mul=True,
+-                    override_orig_dtype=torch.float32,
+-                    fp32_residual=True,
++                    fp32_residual=False,
+                 )
+                 if get_global_server_args().rl_on_policy_target is not None
+                 else {}
+diff --git a/python/sglang/srt/models/qwen2_moe.py b/python/sglang/srt/models/qwen2_moe.py
+index bbb883a2d..9bad2d1e0 100644
+--- a/python/sglang/srt/models/qwen2_moe.py
++++ b/python/sglang/srt/models/qwen2_moe.py
+@@ -596,7 +596,17 @@ class Qwen2MoeModel(nn.Module):
+             prefix=add_prefix("layers", prefix),
+         )
+         if self.pp_group.is_last_rank:
+-            self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
++            norm_kwargs = (
++                dict(
++                    cast_x_before_out_mul=True,
++                    fp32_residual=False,
++                )
++                if get_global_server_args().rl_on_policy_target is not None
++                else {}
++            )
++            self.norm = RMSNorm(
++                config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
++            )
+         else:
+             self.norm = PPMissingLayer(return_tuple=True)
+ 
+diff --git a/python/sglang/srt/models/qwen3.py b/python/sglang/srt/models/qwen3.py
+index b056317e4..351b18684 100644
+--- a/python/sglang/srt/models/qwen3.py
++++ b/python/sglang/srt/models/qwen3.py
+@@ -90,8 +90,8 @@ class Qwen3Attention(nn.Module):
+ 
+         norm_kwargs = (
+             dict(
+-                weight_dtype=torch.float32,
+                 cast_x_before_out_mul=True,
++                fp32_residual=False,
+             )
+             if get_global_server_args().rl_on_policy_target is not None
+             else {}
+@@ -242,10 +242,8 @@ class Qwen3DecoderLayer(nn.Module):
+ 
+         norm_kwargs = (
+             dict(
+-                weight_dtype=torch.float32,
+                 cast_x_before_out_mul=True,
+-                override_orig_dtype=torch.float32,
+-                fp32_residual=True,
++                fp32_residual=False,
+             )
+             if get_global_server_args().rl_on_policy_target is not None
+             else {}
 diff --git a/python/sglang/srt/models/qwen3_moe.py b/python/sglang/srt/models/qwen3_moe.py
-index 912891b6a7..fd67a7b580 100644
+index 3fcf0cfa0..bfef7bcf8 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)
+@@ -22,6 +22,7 @@ import math
+ from typing import Any, Dict, Iterable, List, Optional, Tuple, TypeVar
  
--        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)
+ import torch
++import torch.nn.functional as F
+ from torch import nn
+ from transformers import PretrainedConfig
  
+@@ -50,7 +51,7 @@ from sglang.srt.layers.moe import (
+ )
+ from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class
+ from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE
+-from sglang.srt.layers.moe.topk import TopK
++from sglang.srt.layers.moe.topk import StandardTopKOutput, TopK
+ from sglang.srt.layers.moe.utils import (
+     RoutingMethodType,
+     filter_moe_weight_param_global_expert,
+@@ -233,6 +234,7 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
+             use_grouped_topk=False,
+             layer_id=layer_id,
+         )
++        self.top_k = config.num_experts_per_tok
+ 
+         self.experts = get_moe_impl_class(quant_config)(
+             num_experts=config.num_experts
+@@ -301,7 +303,22 @@ class Qwen3MoeSparseMoeBlock(nn.Module):
+ 
+         # router_logits: (num_tokens, n_experts)
+         router_logits, _ = self.gate(hidden_states)
+-        topk_output = self.topk(hidden_states, router_logits)
++
++        if get_global_server_args().rl_on_policy_target is not None:
++            routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
++            routing_weights, selected_experts = torch.topk(
++                routing_weights, self.top_k, dim=-1
++            )
++            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
++            routing_weights = routing_weights.to(hidden_states.dtype)
++            topk_output = StandardTopKOutput(
++                topk_weights=routing_weights,
++                topk_ids=selected_experts,
++                router_logits=router_logits,
++            )
++        else:
++            topk_output = self.topk(hidden_states, router_logits)
++
+         final_hidden_states = self.experts(hidden_states, topk_output)
          if (
+             self.tp_size > 1
+@@ -482,13 +499,14 @@ class Qwen3MoeAttention(nn.Module):
+         )
+         self.compatible_with_fused_kv_buffer = (
+             False if isinstance(self.rotary_emb, MRotaryEmbedding) else True
+-        )
++        ) and (get_global_server_args().rl_on_policy_target is None)
+         self.compatible_with_fused_qk_norm_rope = (
+             not isinstance(self.rotary_emb, MRotaryEmbedding)
+         ) and self.head_dim in (64, 128, 256)
+         self.use_fused_qk_norm_rope = (
+             get_global_server_args().enable_fused_qk_norm_rope
+             and self.compatible_with_fused_qk_norm_rope
++            and (get_global_server_args().rl_on_policy_target is None)
+         )
+         self._used_fused_qk_norm_rope_last_call = False
+ 
+@@ -501,8 +519,16 @@ class Qwen3MoeAttention(nn.Module):
+             prefix=add_prefix("attn", prefix),
+         )
+ 
+-        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
+-        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps)
++        norm_kwargs = (
++            dict(
++                cast_x_before_out_mul=True,
++                fp32_residual=False,
++            )
++            if get_global_server_args().rl_on_policy_target is not None
++            else {}
++        )
++        self.q_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
++        self.k_norm = RMSNorm(self.head_dim, eps=rms_norm_eps, **norm_kwargs)
+         self.alt_stream = alt_stream
+ 
+     def op_prepare(self, state):
+@@ -743,9 +769,19 @@ class Qwen3MoeDecoderLayer(nn.Module):
+                 quant_config=quant_config,
+                 prefix=add_prefix("mlp", prefix),
+             )
+-        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
++        norm_kwargs = (
++            dict(
++                cast_x_before_out_mul=True,
++                fp32_residual=False,
++            )
++            if get_global_server_args().rl_on_policy_target is not None
++            else {}
++        )
++        self.input_layernorm = RMSNorm(
++            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
++        )
+         self.post_attention_layernorm = RMSNorm(
+-            config.hidden_size, eps=config.rms_norm_eps
++            config.hidden_size, eps=config.rms_norm_eps, **norm_kwargs
+         )
+ 
+         self.layer_communicator = LayerCommunicator(
 diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py
-index 7746b24459..57b65fe06f 100644
+index d641826e3..3abc39ef3 100644
 --- a/python/sglang/srt/models/qwen3_vl.py
 +++ b/python/sglang/srt/models/qwen3_vl.py
-@@ -1005,14 +1005,19 @@ class Qwen3LLMModel(Qwen3Model):
+@@ -711,14 +711,19 @@ class Qwen3LLMModel(Qwen3Model):
                      hidden_states + residual if residual is not None else hidden_states
                  )
  
@@ -1835,434 +2783,139 @@ index 7746b24459..57b65fe06f 100644
              hidden_states, residual = layer(
                  positions,
                  hidden_states,
-diff --git a/python/sglang/srt/multimodal/processors/glm4v.py b/python/sglang/srt/multimodal/processors/glm4v.py
-index a44f14b6ca..6d6c65ea49 100644
---- a/python/sglang/srt/multimodal/processors/glm4v.py
-+++ b/python/sglang/srt/multimodal/processors/glm4v.py
-@@ -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 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 (
-@@ -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
-+        self.IM_START_TOKEN_ID = self.IMAGE_START_TOKEN_ID
-+        self.IM_END_TOKEN_ID = self.IMAGE_END_TOKEN_ID
- 
-         # Vision config
-         self.IMAGE_FACTOR = 28
-@@ -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=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=image_embeddings,
-+            )
-+        ]
-+
-+        mrope_positions, mrope_position_delta = MRotaryEmbedding.get_rope_index_glm4v(
-+            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,
-+        )
-+
-+        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 3f102567d0..6fb3899021 100644
---- a/python/sglang/srt/multimodal/processors/qwen_vl.py
-+++ b/python/sglang/srt/multimodal/processors/qwen_vl.py
-@@ -499,7 +499,7 @@ class QwenVLImageProcessor(SGLangBaseProcessor):
-         **kwargs,
-     ):
-         entry_time = time.perf_counter()
--        base_output = self.load_mm_data(
-+        base_output = self.legacy_load_mm_data(
-             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 d91ced805f..4c8774bb64 100644
+index b080aeb16..957a613fa 100644
 --- a/python/sglang/srt/server_args.py
 +++ b/python/sglang/srt/server_args.py
-@@ -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
-     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.",
+@@ -580,6 +580,7 @@ class ServerArgs:
+     cuda_graph_max_bs: Optional[int] = None
+     cuda_graph_bs: Optional[List[int]] = None
+     disable_cuda_graph: bool = False
++    disable_draft_cuda_graph: bool = False
+     disable_cuda_graph_padding: bool = False
+     enable_profile_cuda_graph: bool = False
+     enable_cudagraph_gc: bool = False
+@@ -2089,7 +2090,16 @@ class ServerArgs:
+             assert (
+                 self.tp_size % (self.dp_size * self.attn_cp_size) == 0
+             ), "tp_size must be divisible by dp_size * attn_cp_size"
+-            assert self.pp_size == 1, "PP is not supported with context parallelism"
++            if self.pp_size > 1:
++                assert (
++                    self.disaggregation_mode == "prefill"
++                    and self.enable_nsa_prefill_context_parallel
++                    and self.nsa_prefill_cp_mode == "round-robin-split"
++                ), (
++                    "PP with context parallelism is only supported for PD prefill "
++                    "with --enable-nsa-prefill-context-parallel and "
++                    "--nsa-prefill-cp-mode round-robin-split."
++                )
+ 
+         if self.moe_dp_size > 1:
+             # The tp_size is the world size, not the real tensor parallel size
+@@ -4491,6 +4501,11 @@ class ServerArgs:
+             action="store_true",
+             help="Disable cuda graph.",
          )
 +        parser.add_argument(
-+            "--disable-indexer-rope-neox-style",
++            "--disable-draft-cuda-graph",
 +            action="store_true",
-+            help="Disable NSA indexer RoPE neox style (equivalent to INDEXER_ROPE_NEOX_STYLE=0). "
-+            "If the environment variable INDEXER_ROPE_NEOX_STYLE is also set and conflicts, an error is raised.",
++            help="Disable cuda graph for draft model in speculative decoding.",
 +        )
          parser.add_argument(
-             "--enable-prefill-context-parallel",
+             "--disable-cuda-graph-padding",
              action="store_true",
+@@ -5636,6 +5651,54 @@ class PortArgs:
+             )
+ 
+         if not server_args.enable_dp_attention:
++            # In multi-node prefill PD with PP/CP, use TCP transport for tokenizer<->scheduler
++            # IPC occasionally stalls in this topology.
++            if (
++                server_args.nnodes > 1
++                and server_args.disaggregation_mode == "prefill"
++                and server_args.dist_init_addr is not None
++            ):
++                if server_args.dist_init_addr.startswith("["):  # ipv6 address
++                    port_num, host = configure_ipv6(server_args.dist_init_addr)
++                    dist_init_addr = (host, str(port_num))
++                else:
++                    dist_init_addr = server_args.dist_init_addr.split(":")
++
++                assert (
++                    len(dist_init_addr) == 2
++                ), "please provide --dist-init-addr as host:port of head node"
++
++                dist_init_host, dist_init_port = dist_init_addr
++                dist_init_port = int(dist_init_port)
++                port_base = dist_init_port + ZMQ_TCP_PORT_DELTA
++                detokenizer_port = port_base + 1
++                rpc_port = port_base + 2
++                metrics_port = port_base + 3
++                scheduler_input_port = port_base + 4
++
++                try:
++                    wait_port_available(dist_init_port, "dist_init_port")
++                    wait_port_available(port_base, "port_base")
++                    wait_port_available(detokenizer_port, "detokenizer_port")
++                    wait_port_available(nccl_port, "nccl_port")
++                    wait_port_available(rpc_port, "rpc_port")
++                    wait_port_available(metrics_port, "metrics_port")
++                    wait_port_available(scheduler_input_port, "scheduler_input_port")
++                except ValueError:
++                    logger.exception(
++                        f"Port is already in use. {dist_init_port=} {port_base=} {detokenizer_port=} {nccl_port=} {scheduler_input_port=}"
++                    )
++                    raise
++
++                return PortArgs(
++                    tokenizer_ipc_name=f"tcp://{dist_init_host}:{port_base}",
++                    scheduler_input_ipc_name=f"tcp://{dist_init_host}:{scheduler_input_port}",
++                    detokenizer_ipc_name=f"tcp://{dist_init_host}:{detokenizer_port}",
++                    nccl_port=nccl_port,
++                    rpc_ipc_name=f"tcp://{dist_init_host}:{rpc_port}",
++                    metrics_ipc_name=f"tcp://{dist_init_host}:{metrics_port}",
++                    tokenizer_worker_ipc_name=tokenizer_worker_ipc_name,
++                )
+             # Normal case, use IPC within a single node
+             return PortArgs(
+                 tokenizer_ipc_name=f"ipc://{tempfile.NamedTemporaryFile(delete=False).name}",
 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 40e859b2d6..2604ae037c 100644
+index 5fe45086c..b283d2e9b 100644
 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
 +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
-@@ -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_()
- 
+@@ -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_()
          num_tokens = bs * self.num_tokens_per_bs
  
-@@ -386,8 +390,12 @@ class EAGLEDraftCudaGraphRunner:
+         # Common inputs
+@@ -350,8 +353,12 @@ class EAGLEDraftCudaGraphRunner:
              forward_batch.out_cache_loc
          )
-         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_(
+         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_(
 +            forward_batch.spec_info.topk_index.clamp(
 +                0, self.model_runner.model_config.vocab_size - 1
 +            )
 +        )
-         buffers.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
-         buffers.req_pool_indices[:raw_bs].copy_(forward_batch.req_pool_indices)
+         self.hidden_states[:raw_bs].copy_(forward_batch.spec_info.hidden_states)
+         self.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 dbb91f555e..a04caefc34 100644
+index ac629c7ee..c039d2350 100644
 --- a/python/sglang/srt/speculative/eagle_info.py
 +++ b/python/sglang/srt/speculative/eagle_info.py
-@@ -776,6 +776,10 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
+@@ -774,6 +774,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)]
@@ -2273,7 +2926,7 @@ index dbb91f555e..a04caefc34 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]
-@@ -807,6 +811,27 @@ class EagleDraftInput(SpecInput, EagleDraftInputV2Mixin):
+@@ -805,6 +809,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])
@@ -2301,20 +2954,37 @@ index dbb91f555e..a04caefc34 100644
  
  
  @dataclass
+diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py
+index 32b3a520a..d7f940147 100644
+--- a/python/sglang/srt/speculative/eagle_worker.py
++++ b/python/sglang/srt/speculative/eagle_worker.py
+@@ -234,7 +234,10 @@ class EAGLEWorker(TpModelWorker):
+         self.cuda_graph_runner = None
+         self.cuda_graph_runner_for_draft_extend = None
+ 
+-        if self.server_args.disable_cuda_graph:
++        if (
++            self.server_args.disable_cuda_graph
++            or self.server_args.disable_draft_cuda_graph
++        ):
+             return
+ 
+         Device2DraftCudaGraphRunner = {
 diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py
-index b0be70d751..44a78d684e 100644
+index 4636128fa..a9b61df39 100644
 --- a/python/sglang/srt/utils/common.py
 +++ b/python/sglang/srt/utils/common.py
-@@ -2157,6 +2157,7 @@ class SafeUnpickler(pickle.Unpickler):
+@@ -2359,6 +2359,8 @@ class SafeUnpickler(pickle.Unpickler):
+         "sglang.srt.model_executor.model_runner.",
          "sglang.srt.layers.",
          "sglang.srt.utils.",
-         "torch_npu.",
++        # --- slime ---
 +        "slime.",
      }
  
      DENY_CLASSES = {
 diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py
-index 3be16446e0..1b2371c839 100644
+index 3be16446e..1b2371c83 100644
 --- a/python/sglang/srt/utils/weight_checker.py
 +++ b/python/sglang/srt/utils/weight_checker.py
 @@ -69,6 +69,9 @@ def _check_tensors(
diff --git a/docker/patch/v0.5.9/megatron.patch b/docker/patch/v0.5.9/megatron.patch
deleted file mode 100644
index 2e6ae436b1..0000000000
--- a/docker/patch/v0.5.9/megatron.patch
+++ /dev/null
@@ -1,795 +0,0 @@
-diff --git a/megatron/core/dist_checkpointing/strategies/common.py b/megatron/core/dist_checkpointing/strategies/common.py
-index 41c21d93d..ef80f72d6 100644
---- a/megatron/core/dist_checkpointing/strategies/common.py
-+++ b/megatron/core/dist_checkpointing/strategies/common.py
-@@ -86,7 +86,7 @@ class TorchCommonLoadStrategy(LoadCommonStrategy):
-                 msc = MultiStorageClientFeature.import_package()
-                 return msc.torch.load(load_path, map_location='cpu')
-             else:
--                return torch.load(load_path, map_location='cpu')
-+                return torch.load(load_path, map_location='cpu', weights_only=False)
-         except FileNotFoundError as e:
-             err_msg = f'Common file {load_path} does not exist'
-             if MultiStorageClientFeature.is_enabled():
-diff --git a/megatron/core/dist_checkpointing/strategies/torch.py b/megatron/core/dist_checkpointing/strategies/torch.py
-index 5a1ea308d..aa701237f 100644
---- a/megatron/core/dist_checkpointing/strategies/torch.py
-+++ b/megatron/core/dist_checkpointing/strategies/torch.py
-@@ -597,10 +597,12 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
-     def _validate_global_shapes(self, metadata, sharded_tensors):
-         for sh_ten in sharded_tensors:
-             if sh_ten.key not in metadata.state_dict_metadata:
--                raise KeyError(
--                    f"{sh_ten.key} from model not in state dict:"
--                    f" {sorted(metadata.state_dict_metadata.keys())}"
--                )
-+                # raise KeyError(
-+                #     f"{sh_ten.key} from model not in state dict:"
-+                #     f" {sorted(metadata.state_dict_metadata.keys())}"
-+                # )
-+                print(f"{sh_ten.key} from model not in state dict, will skip")
-+                continue
-             loaded_shape = metadata.state_dict_metadata[sh_ten.key].size
-             expected_shape = self._expected_shape(sh_ten)
-             if loaded_shape != expected_shape:
-@@ -630,7 +632,7 @@ class MCoreLoadPlanner(DefaultLoadPlanner):
-         tensor_metadata = self.metadata.state_dict_metadata
-         metadata_with_sizes = [
-             (tensor_metadata[key], tensor_metadata[key].size, sharded_tensor)
--            for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items()
-+            for key, sharded_tensor in self.allow_shape_mismatch_sharded_tensors.items() if key in tensor_metadata
-         ]
-         try:
-             # Temporarily set sizes to expected shapes
-@@ -959,6 +961,7 @@ class TorchDistLoadShardedStrategy(LoadShardedStrategy):
-             planner=MCoreLoadPlanner(
-                 shapes_validation_sharded_tensors=flexible_shape_sharded_tensors,
-                 allow_shape_mismatch_sharded_tensors=allow_shape_mismatch_sharded_tensors,
-+                allow_partial_load=True,
-             ),
-         )
- 
-diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py
-index acb93ef78..d239db4ab 100644
---- a/megatron/core/extensions/transformer_engine.py
-+++ b/megatron/core/extensions/transformer_engine.py
-@@ -408,6 +408,7 @@ class TELinear(te.pytorch.Linear):
-         )
- 
-         for param in self.parameters():
-+            setattr(param, "parallel_mode", parallel_mode)
-             if is_expert:
-                 # Reduce the gradient on the expert_data_parallel group for expert linear layers
-                 setattr(param, "allreduce", not self.expert_parallel)
-@@ -1161,6 +1162,61 @@ class TEDotProductAttention(te.pytorch.DotProductAttention):
- 
- 
- if HAVE_TE and is_te_min_version("1.9.0.dev0"):
-+    def ceil_div(x: int, y: int) -> int:
-+        return (x + y - 1) // y
-+
-+    class _FakeInt4QuantizationSTE(torch.autograd.Function):
-+        @staticmethod
-+        def forward(ctx, x, group_size):
-+            m, n = x.shape
-+            block_size_m, block_size_n = 1, group_size
-+
-+
-+            m_padded = ceil_div(m, block_size_m) * block_size_m
-+            n_padded = ceil_div(n, block_size_n) * block_size_n
-+
-+            x_padded = torch.zeros(
-+                (m_padded, n_padded),
-+                dtype=x.dtype, device=x.device
-+            )
-+            x_padded[:m, :n] = x
-+
-+            x_view = x_padded.view(
-+                m_padded // block_size_m,
-+                block_size_m,
-+                n_padded // block_size_n,
-+                block_size_n
-+            )
-+
-+            x_max = x_view.abs().float().amax(dim=(1, 3), keepdim=True)
-+            q_max = 7
-+            x_scale = x_max / q_max
-+
-+            x_scale = x_scale.clamp(min=1e-5)
-+
-+            x_div = x_view / x_scale
-+            x_round = torch.round(x_div)
-+
-+            x_q_clamped = x_round.clamp(-q_max, q_max)
-+
-+            x_dequant_view = x_q_clamped * x_scale
-+
-+            x_dequant_full = x_dequant_view.view_as(x_padded)
-+            x_out = x_dequant_full[:m, :n].contiguous().to(x.dtype)
-+
-+            return x_out
-+
-+        @staticmethod
-+        def backward(ctx, grad_output):
-+            return grad_output, None
-+
-+    def fake_int4_quantization_ste(x, group_size):
-+        x_out = _FakeInt4QuantizationSTE.apply(x, group_size)
-+        
-+        if hasattr(x, 'main_grad'):
-+            x_out.main_grad = x.main_grad
-+            
-+        return x_out
- 
-     class TEGroupedLinear(te.pytorch.GroupedLinear):
-         """
-@@ -1351,6 +1407,7 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"):
-             _is_first_microbatch = (
-                 None if self.disable_parameter_transpose_cache else self.is_first_microbatch
-             )
-+
-             out = super().forward(x, m_splits, is_first_microbatch=_is_first_microbatch)
-             self.is_first_microbatch = False
- 
-@@ -1361,6 +1418,20 @@ if HAVE_TE and is_te_min_version("1.9.0.dev0"):
-                 return out
-             return out, None
- 
-+        def _get_weight_tensors(self):
-+            """Get the weight tensors of the module."""
-+            weight_tensors = super()._get_weight_tensors()
-+
-+            if os.getenv("OPEN_TRAINING_INT4_FAKE_QAT_FLAG", "0") == "1":
-+                group_size = int(os.getenv("OPEN_TRAINING_INT4_GROUP_SIZE", "128"))
-+
-+                weight_tensors = [
-+                    fake_int4_quantization_ste(w, group_size) 
-+                    for w in weight_tensors
-+                ]
-+                
-+            return weight_tensors
-+
-         def _encode_extra_state(self, state):
-             # TE 2.0 changed the format of extra_state to be a byte tensor
-             if is_te_min_version("2.0.0"):
-diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py
-index 1fd5dcfae..c9aeef1f0 100644
---- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py
-+++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py
-@@ -385,6 +385,7 @@ def rotary_fwd_kv_kernel(
-     SIN,
-     emb_dim: tl.constexpr,
-     k_dim: tl.constexpr,
-+    k_dim_ceil: tl.constexpr,
-     v_dim: tl.constexpr,
-     head_num: tl.constexpr,
-     batch_size,
-@@ -434,21 +435,27 @@ def rotary_fwd_kv_kernel(
-     cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2))
-     sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2))
- 
--    KV_ptr = KV + pid_m * stride_kv_seq + pid_head * BLOCK_H * stride_kv_nheads
--    kv_off = tl.arange(0, BLOCK_H)[:, None] * stride_kv_nheads
--    mask = kv_off < head_num * stride_kv_nheads
--    k_in_off = kv_off + tl.arange(0, k_dim)[None, :]
--    v_in_off = kv_off + k_dim + tl.arange(0, v_dim)[None, :]
--    k = tl.load(KV_ptr + k_in_off, mask=mask)
--    v = tl.load(KV_ptr + v_in_off, mask=mask)
-+    KV_ptr = KV + pid_m * stride_kv_seq # + pid_head * BLOCK_H * stride_kv_nheads
-+    ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H
-+    kj_range = tl.arange(0, k_dim_ceil)[None, :]
-+    mask_k = (ki_range < head_num) & (kj_range < k_dim)
-+    mask_v = ki_range < head_num
-+    k_off = ki_range * stride_kv_nheads + kj_range
-+    if v_dim > 0:
-+        v_off = ki_range * stride_kv_nheads + k_dim + tl.arange(0, v_dim)[None, :]
-+        v = tl.load(KV_ptr + v_off, mask=mask_v)
-+    else:
-+        v = tl.zeros((BLOCK_H, 1), dtype=KV.dtype.element_ty)
-+    k = tl.load(KV_ptr + k_off, mask=mask_k)
- 
--    K_ptr = O_KEY + pid_m * stride_k_seq + pid_head * BLOCK_H * stride_k_nheads
--    V_ptr = O_VALUE + pid_m * stride_v_seq + pid_head * BLOCK_H * stride_v_nheads
-+    K_ptr = O_KEY + pid_m * stride_k_seq # + pid_head * BLOCK_H * stride_k_nheads
-+    V_ptr = O_VALUE + pid_m * stride_v_seq # + pid_head * BLOCK_H * stride_v_nheads
- 
--    k_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + tl.arange(0, k_dim)[None, :]
--    v_out_off = tl.arange(0, BLOCK_H)[:, None] * stride_v_nheads + tl.arange(0, v_dim)[None, :]
--    tl.store(K_ptr + k_out_off, k, mask=mask)
--    tl.store(V_ptr + v_out_off, v, mask=mask)
-+    k_out_off = ki_range * stride_k_nheads + kj_range
-+    tl.store(K_ptr + k_out_off, k, mask=mask_k)
-+    if v_dim > 0:
-+        v_out_off = ki_range * stride_v_nheads + tl.arange(0, v_dim)[None, :]
-+        tl.store(V_ptr + v_out_off, v, mask=mask_v)
- 
-     EMB = K_POS_EMB + pid_m * stride_emb_seq
-     # x1 = t[..., 0::2], x2 = t[..., 1::2]
-@@ -460,14 +467,16 @@ def rotary_fwd_kv_kernel(
-     x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2)
-     x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2)
- 
-+    x_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H
-+    mask_x = x_range < head_num
-     x_left_off = (
--        tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads
-+        x_range * stride_k_nheads
-         + k_dim
-         + tl.arange(0, emb_dim // 2)[None, :]
-     )
-     x_right_off = x_left_off + emb_dim // 2
--    tl.store(K_ptr + x_left_off, x_left, mask=mask)
--    tl.store(K_ptr + x_right_off, x_right, mask=mask)
-+    tl.store(K_ptr + x_left_off, x_left, mask=mask_x)
-+    tl.store(K_ptr + x_right_off, x_right, mask=mask_x)
- 
- 
- @triton.autotune(
-@@ -493,6 +502,7 @@ def rotary_bwd_kv_kernel(
-     SIN,
-     emb_dim: tl.constexpr,
-     k_dim: tl.constexpr,
-+    k_dim_ceil: tl.constexpr,
-     v_dim: tl.constexpr,
-     head_num: tl.constexpr,
-     batch_size,
-@@ -533,27 +543,32 @@ def rotary_bwd_kv_kernel(
-     else:
-         token_idx = _get_thd_token_idx(cu_seqlens_kv, pid_m, seq_num, cp_rank, cp_size)
- 
--    dKV_ptr = dKV + pid_m * stride_dkv_seq + pid_head * BLOCK_H * stride_dkv_nheads
--    dkv_off = tl.arange(0, BLOCK_H)[:, None] * stride_dkv_nheads
--    mask = dkv_off < head_num * stride_dkv_nheads
--    dk_out_off = dkv_off + tl.arange(0, k_dim)[None, :]
--    dv_out_off = dkv_off + k_dim + tl.arange(0, v_dim)[None, :]
--
--    dK_ptr = dK + pid_m * stride_dk_seq + pid_head * BLOCK_H * stride_dk_nheads
--    dV_ptr = dV + pid_m * stride_dv_seq + pid_head * BLOCK_H * stride_dv_nheads
--    dk_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + tl.arange(0, k_dim)[None, :]
--    dv_in_off = tl.arange(0, BLOCK_H)[:, None] * stride_dv_nheads + tl.arange(0, v_dim)[None, :]
--    dk = tl.load(dK_ptr + dk_in_off, mask=mask)
--    dv = tl.load(dV_ptr + dv_in_off, mask=mask)
--    tl.store(dKV_ptr + dk_out_off, dk, mask=mask)
--    tl.store(dKV_ptr + dv_out_off, dv, mask=mask)
-+    dKV_ptr = dKV + pid_m * stride_dkv_seq # + pid_head * BLOCK_H * stride_dkv_nheads
-+    ki_range = tl.arange(0, BLOCK_H)[:, None] + pid_head * BLOCK_H 
-+    kj_range = tl.arange(0, k_dim_ceil)[None, :]
-+    mask_k = (ki_range < head_num) & (kj_range < k_dim)
-+    mask_v = ki_range < head_num
-+    dk_out_off = ki_range * stride_dkv_nheads + kj_range
-+
-+    dK_ptr = dK + pid_m * stride_dk_seq # + pid_head * BLOCK_H * stride_dk_nheads
-+    dV_ptr = dV + pid_m * stride_dv_seq # + pid_head * BLOCK_H * stride_dv_nheads
-+    dk_in_off = ki_range * stride_dk_nheads + kj_range
-+
-+    dk = tl.load(dK_ptr + dk_in_off, mask=mask_k)
-+    tl.store(dKV_ptr + dk_out_off, dk, mask=mask_k)
-+    
-+    if v_dim > 0:
-+        dv_out_off = ki_range * stride_dkv_nheads + k_dim + tl.arange(0, v_dim)[None, :]
-+        dv_in_off = ki_range * stride_dv_nheads + tl.arange(0, v_dim)[None, :]
-+        dv = tl.load(dV_ptr + dv_in_off, mask=mask_v)
-+        tl.store(dKV_ptr + dv_out_off, dv, mask=mask_v)
- 
-     if pid_head == 0:
-         x_left_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32)
-         x_right_accum = tl.zeros((BLOCK_H, emb_dim // 2), dtype=tl.float32)
-         for i in tl.static_range(triton.cdiv(head_num, BLOCK_H)):
--            dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads
--            x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim
-+            dK_ptr = dK + pid_m * stride_dk_seq # + i * BLOCK_H * stride_dk_nheads
-+            x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim + i * BLOCK_H * stride_dk_nheads
-             mask = x_off < head_num * stride_dk_nheads
-             x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :]
-             x_right_off = x_left_off + emb_dim // 2
-@@ -632,6 +647,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
- 
-         o_key = kv.new_empty(total_seqlen, nheads, emb_dim + k_dim)
-         o_value = kv.new_empty(total_seqlen, nheads, v_dim)
-+        k_dim_ceil = triton.next_power_of_2(k_dim)
- 
-         grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"]))
-         rotary_fwd_kv_kernel[grid](
-@@ -643,6 +659,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
-             sin,
-             emb_dim,
-             k_dim,
-+            k_dim_ceil,
-             v_dim,
-             nheads,
-             batch_size,
-@@ -700,6 +717,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
- 
-         d_kv = dk.new_empty(total_seqlen, nheads, ctx.k_dim + ctx.v_dim)
-         d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim)
-+        k_dim_ceil = triton.next_power_of_2(ctx.k_dim)
- 
-         grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"]))
-         rotary_bwd_kv_kernel[grid](
-@@ -711,6 +729,7 @@ class ApplyMLARotaryEmbKV(torch.autograd.Function):
-             sin,
-             ctx.emb_dim,
-             ctx.k_dim,
-+            k_dim_ceil,
-             ctx.v_dim,
-             nheads,
-             batch_size,
-diff --git a/megatron/core/models/common/embeddings/rotary_pos_embedding.py b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-index 5d7b69cd3..2e0a26815 100644
---- a/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-+++ b/megatron/core/models/common/embeddings/rotary_pos_embedding.py
-@@ -348,6 +348,7 @@ class MultimodalRotaryEmbedding(nn.Module):
- 
-         # shape (seq_length, bs, 1, 2 * dim)
-         emb = emb[..., None, :].transpose(0, 1).contiguous()
-+        packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd'
-         if packed_seq_params is not None and packed_seq_params.local_cp_size is not None:
-             if packed_seq_params.local_cp_size > 1:
-                 # Set CP group to dynamic CP group for CP slicing
-@@ -357,7 +358,9 @@ class MultimodalRotaryEmbedding(nn.Module):
-                 cp_group = None
-         else:
-             cp_group = self.cp_group
--        if cp_group is not None and cp_group.size() > 1:
-+        # For THD (packed sequence) format, skip CP slicing here — it is handled
-+        # per-sequence inside _apply_rotary_pos_emb_thd instead (same as RotaryEmbedding).
-+        if cp_group is not None and cp_group.size() > 1 and not packed_seq:
-             # slice rotary_pos_emb along sequence dimension and select the parition of the current
-             # CP rank
-             emb = get_pos_emb_on_this_cp_rank(emb, 0, cp_group)
-diff --git a/megatron/core/models/common/language_module/language_module.py b/megatron/core/models/common/language_module/language_module.py
-index 13d74aa52..060898a7a 100644
---- a/megatron/core/models/common/language_module/language_module.py
-+++ b/megatron/core/models/common/language_module/language_module.py
-@@ -184,7 +184,15 @@ class LanguageModule(MegatronModule):
-             assert (
-                 column_parallel_linear is not None
-             ), "column_parallel_linear cannot be None when not using fused linear cross entropy."
--            logits, _ = column_parallel_linear(hidden, **col_linear_kwargs)
-+            # output
-+            output_layer_params = {k: v.detach() for k, v in column_parallel_linear.named_parameters()}
-+            output_layer_buffers = dict(column_parallel_linear.named_buffers())
-+            logits, _ = torch.func.functional_call(
-+                column_parallel_linear,
-+                {**output_layer_params, **output_layer_buffers},
-+                (hidden,),
-+                col_linear_kwargs,
-+            )
- 
-             return self.compute_language_model_loss(labels, logits)
- 
-diff --git a/megatron/core/models/gpt/gpt_layer_specs.py b/megatron/core/models/gpt/gpt_layer_specs.py
-index e21127b87..712793853 100755
---- a/megatron/core/models/gpt/gpt_layer_specs.py
-+++ b/megatron/core/models/gpt/gpt_layer_specs.py
-@@ -188,6 +188,8 @@ def get_gpt_layer_with_transformer_engine_spec(
-     use_kitchen: bool = False,
-     use_te_activation_func: bool = False,
-     fallback_to_eager_attn: bool = False,
-+    post_self_attn_layernorm: bool = False,
-+    post_mlp_layernorm: bool = False,
- ) -> ModuleSpec:
-     """Use this spec to use lower-level Transformer Engine modules (required for fp8 training).
- 
-@@ -260,6 +262,8 @@ def get_gpt_layer_with_transformer_engine_spec(
-         mlp=mlp,
-         sharded_state_dict_keys_map=sharded_state_dict_keys_map,
-         normalization=normalization,
-+        post_self_attn_layernorm=post_self_attn_layernorm,
-+        post_mlp_layernorm=post_mlp_layernorm,
-     )
- 
- 
-@@ -349,6 +353,8 @@ def get_transformer_layer_spec_for_backend(
-     mlp: ModuleSpec,
-     sharded_state_dict_keys_map: Optional[dict] = None,
-     normalization: Optional[str] = None,
-+    post_self_attn_layernorm: bool = False,
-+    post_mlp_layernorm: bool = False,
- ) -> ModuleSpec:
-     """Helper function to get module spec for TransformerLayer"""
- 
-@@ -371,9 +377,11 @@ def get_transformer_layer_spec_for_backend(
-             input_layernorm=input_layernorm,
-             self_attention=attention,
-             self_attn_bda=get_bias_dropout_add,
-+            post_self_attn_layernorm=TENorm if post_self_attn_layernorm else IdentityOp,
-             pre_mlp_layernorm=pre_mlp_layernorm,
-             mlp=mlp,
-             mlp_bda=get_bias_dropout_add,
-+            post_mlp_layernorm=TENorm if post_mlp_layernorm else IdentityOp,
-             sharded_state_dict_keys_map=sharded_state_dict_keys_map,
-         ),
-     )
-diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py
-index a1230568c..1fd52f65a 100644
---- a/megatron/core/models/gpt/gpt_model.py
-+++ b/megatron/core/models/gpt/gpt_model.py
-@@ -446,6 +446,7 @@ class GPTModel(LanguageModule):
-         *,
-         inference_params: Optional[BaseInferenceContext] = None,
-         loss_mask: Optional[Tensor] = None,
-+        mtp_kwargs: Optional[dict] = {},
-     ) -> Tensor:
-         """Forward function of the GPT Model This function passes the input tensors
-         through the embedding layer, and then the decoder and finally into the post
-@@ -508,6 +509,7 @@ class GPTModel(LanguageModule):
-             runtime_gather_output=runtime_gather_output,
-             extra_block_kwargs=extra_block_kwargs,
-             inference_context=inference_context,
-+            mtp_kwargs=mtp_kwargs,
-         )
- 
-     def _postprocess(
-@@ -529,6 +531,7 @@ class GPTModel(LanguageModule):
-         runtime_gather_output=None,
-         extra_block_kwargs=None,
-         inference_context=None,
-+        mtp_kwargs={},
-     ):
-         """Postprocesses decoder hidden states to generate logits or compute loss.
- 
-@@ -543,7 +546,8 @@ class GPTModel(LanguageModule):
-         output_weight = None
-         if self.share_embeddings_and_output_weights:
-             output_weight = self.shared_embedding_or_output_weight()
--        if mtp_in_postprocess:
-+
-+        if mtp_in_postprocess and mtp_kwargs.get('mtp_labels', None) is not None:
-             hidden_states = self.mtp(
-                 input_ids=input_ids,
-                 position_ids=position_ids,
-@@ -563,13 +567,18 @@ class GPTModel(LanguageModule):
-             return hidden_states
- 
-         # Skip when mtp_num_layers is None or 0
--        if self.config.mtp_num_layers:
--            mtp_labels = labels.clone()
-+        if self.config.mtp_num_layers and mtp_kwargs.get('mtp_labels', None) is not None:
-+            mtp_labels = mtp_kwargs['mtp_labels'].clone()
-+            mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
-+
-             hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)
-             hidden_states = hidden_states_list[0]
-             if loss_mask is None:
-                 # if loss_mask is not provided, use all ones as loss_mask
-                 loss_mask = torch.ones_like(mtp_labels)
-+            else:
-+                # Otherwise, roll the loss_mask to keep up with the mtp_labels
-+                loss_mask, _ = roll_tensor(loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group, packed_seq_params=packed_seq_params)
-             for mtp_layer_number in range(self.config.mtp_num_layers):
-                 # Calc loss for the current Multi-Token Prediction (MTP) layers.
-                 mtp_labels, _ = roll_tensor(
-@@ -595,7 +604,7 @@ class GPTModel(LanguageModule):
-                     sequence_parallel_enabled=self.output_layer.sequence_parallel,
-                     column_parallel_linear=self.output_layer,
-                     col_linear_kwargs={
--                        'weight': output_weight,
-+                        'weight': output_weight.detach() if output_weight else None,
-                         'runtime_gather_output': runtime_gather_output,
-                     },
-                 )
-diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py
-index 6e093f96f..eac21a3ea 100644
---- a/megatron/core/optimizer/distrib_optimizer.py
-+++ b/megatron/core/optimizer/distrib_optimizer.py
-@@ -677,6 +677,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
-                 # TE FusedAdam will not accumulate step for empty param groups, so we need to
-                 # align the step across param groups.
-                 param_group["step"] = int(step)
-+            if "step" in param_group and param_group["step"] is None:
-+                del param_group["step"]
- 
-         # Grad scaler state.
-         if self.grad_scaler:
-@@ -1646,6 +1648,8 @@ class DistributedOptimizer(MixedPrecisionOptimizer):
-                             if key == 'padding':
-                                 tensors[key] = LocalNonpersistentObject(tensors[key])
-                                 continue
-+                            if key == 'step':
-+                                continue
-                             assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), (
-                                 tensors[key].shape,
-                                 gbuf_local_start,
-diff --git a/megatron/core/parallel_state.py b/megatron/core/parallel_state.py
-index a273002b9..4f821cfd5 100644
---- a/megatron/core/parallel_state.py
-+++ b/megatron/core/parallel_state.py
-@@ -11,6 +11,7 @@ from typing import Callable, List, Optional
- 
- import numpy as np
- import torch
-+import torch.distributed as dist
- 
- from .utils import GlobalMemoryBuffer, is_torch_min_version
- 
-diff --git a/megatron/core/pipeline_parallel/p2p_communication.py b/megatron/core/pipeline_parallel/p2p_communication.py
-index ac839c21f..f18309217 100644
---- a/megatron/core/pipeline_parallel/p2p_communication.py
-+++ b/megatron/core/pipeline_parallel/p2p_communication.py
-@@ -26,22 +26,22 @@ def _batched_p2p_ops(
-     ops = []
-     if tensor_send_prev is not None:
-         send_prev_op = torch.distributed.P2POp(
--            torch.distributed.isend, tensor_send_prev, prev_pipeline_rank, group
-+            torch.distributed.isend, tensor_send_prev, prev_pipeline_rank,
-         )
-         ops.append(send_prev_op)
-     if tensor_recv_prev is not None:
-         recv_prev_op = torch.distributed.P2POp(
--            torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank, group
-+            torch.distributed.irecv, tensor_recv_prev, prev_pipeline_rank,
-         )
-         ops.append(recv_prev_op)
-     if tensor_send_next is not None:
-         send_next_op = torch.distributed.P2POp(
--            torch.distributed.isend, tensor_send_next, next_pipeline_rank, group
-+            torch.distributed.isend, tensor_send_next, next_pipeline_rank,
-         )
-         ops.append(send_next_op)
-     if tensor_recv_next is not None:
-         recv_next_op = torch.distributed.P2POp(
--            torch.distributed.irecv, tensor_recv_next, next_pipeline_rank, group
-+            torch.distributed.irecv, tensor_recv_next, next_pipeline_rank,
-         )
-         ops.append(recv_next_op)
-     if len(ops) > 0:
-diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py
-index 28cff06f5..58dc4bb70 100644
---- a/megatron/core/transformer/moe/moe_utils.py
-+++ b/megatron/core/transformer/moe/moe_utils.py
-@@ -587,6 +587,9 @@ def topk_routing_with_score_function(
-         else:
-             return torch.topk(scores, k=topk, dim=1)
- 
-+    from slime.utils.routing_replay import get_routing_replay_compute_topk
-+    compute_topk = get_routing_replay_compute_topk(compute_topk)
-+
-     if score_function == "softmax":
-         if use_pre_softmax:
-             scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)
-diff --git a/megatron/core/transformer/moe/router.py b/megatron/core/transformer/moe/router.py
-index 16fc9d9af..517944f25 100644
---- a/megatron/core/transformer/moe/router.py
-+++ b/megatron/core/transformer/moe/router.py
-@@ -201,6 +201,9 @@ class TopKRouter(Router):
-             self.global_tokens_per_expert = None
-             self.ga_steps = None
- 
-+        from slime.utils.routing_replay import register_routing_replay
-+        register_routing_replay(self)
-+
-     def _maintain_float32_expert_bias(self):
-         """
-         Maintain the expert bias in float32.
-diff --git a/megatron/core/transformer/multi_token_prediction.py b/megatron/core/transformer/multi_token_prediction.py
-index a8f4abfcd..f33f6f05e 100755
---- a/megatron/core/transformer/multi_token_prediction.py
-+++ b/megatron/core/transformer/multi_token_prediction.py
-@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Union
- 
- import torch
- from torch import Tensor
-+import warnings
- 
- from megatron.core import InferenceParams, parallel_state, tensor_parallel
- from megatron.core.dist_checkpointing.mapping import ShardedStateDict
-@@ -714,17 +715,19 @@ class MultiTokenPredictionLayer(MegatronModule):
-             cp_group=self.cp_group,
-             packed_seq_params=packed_seq_params,
-         )
--        position_ids, _ = roll_tensor(
--            position_ids,
--            shifts=-1,
--            dims=-1,
--            cp_group=self.cp_group,
--            packed_seq_params=packed_seq_params,
--        )
-+        if position_ids is not None:
-+            position_ids, _ = roll_tensor(
-+                position_ids,
-+                shifts=-1,
-+                dims=-1,
-+                cp_group=self.cp_group,
-+                packed_seq_params=packed_seq_params,
-+            )
-         # embedding
-         decoder_input = embedding(input_ids=input_ids, position_ids=position_ids)
-+        decoder_input = decoder_input.detach()
- 
--        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)
-+        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=False)
- 
-         return input_ids, position_ids, decoder_input, hidden_states
- 
-@@ -826,6 +829,51 @@ class MultiTokenPredictionLayer(MegatronModule):
-         return hidden_states
- 
-     def _checkpointed_forward(self, forward_func, *args, **kwargs):
-+        """Wrap `forward_func` with activation checkpointing while only passing tensors.
-+
-+        Non-tensor arguments (e.g., configuration objects, None) are captured via closure so
-+        that checkpoint implementations never receive them directly, avoiding save_for_backward
-+        issues with non-tensor inputs.
-+        """
-+
-+        # TODO(jiajun): Is there any better implementation here?
-+        positional_specs = []
-+        kw_specs = []
-+        tensor_args: List[torch.Tensor] = []
-+
-+        for arg in args:
-+            if torch.is_tensor(arg):
-+                positional_specs.append(('tensor', len(tensor_args)))
-+                tensor_args.append(arg)
-+            else:
-+                positional_specs.append(('const', arg))
-+
-+        for key, value in kwargs.items():
-+            if torch.is_tensor(value):
-+                kw_specs.append((key, ('tensor', len(tensor_args))))
-+                tensor_args.append(value)
-+            else:
-+                kw_specs.append((key, ('const', value)))
-+
-+        def run(*flat_tensor_args):
-+            rebuilt_args = []
-+            for spec_type, payload in positional_specs:
-+                if spec_type == 'tensor':
-+                    rebuilt_args.append(flat_tensor_args[payload])
-+                else:
-+                    rebuilt_args.append(payload)
-+
-+            rebuilt_kwargs = {}
-+            for key, (spec_type, payload) in kw_specs:
-+                if spec_type == 'tensor':
-+                    rebuilt_kwargs[key] = flat_tensor_args[payload]
-+                else:
-+                    rebuilt_kwargs[key] = payload
-+
-+            return forward_func(*rebuilt_args, **rebuilt_kwargs)
-+
-+        tensor_args_tuple = tuple(tensor_args)
-+
-         def checkpoint_handler():
-             """Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`"""
-             if self.config.fp8:
-@@ -836,12 +884,11 @@ class MultiTokenPredictionLayer(MegatronModule):
-                     self.config.distribute_saved_activations,
-                     tensor_parallel.random.get_cuda_rng_tracker,
-                     parallel_state.get_tensor_model_parallel_group(),
--                    *args,
--                    **kwargs,
-+                    *tensor_args_tuple,
-                 )
-             else:
-                 return tensor_parallel.checkpoint(
--                    forward_func, self.config.distribute_saved_activations, *args, *kwargs.values()
-+                    run, self.config.distribute_saved_activations, *tensor_args_tuple
-                 )
- 
-         if self.config.recompute_method == 'uniform':
-diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py
-index e2705bd9f..a0aa109b5 100644
---- a/megatron/core/transformer/transformer_config.py
-+++ b/megatron/core/transformer/transformer_config.py
-@@ -210,6 +210,9 @@ class TransformerConfig(ModelParallelConfig):
-     attention_output_gate: bool = False
-     """Whether to apply output gate to the attention layers."""
- 
-+    post_self_attn_layernorm: bool = False
-+    post_mlp_layernorm: bool = False
-+
-     test_mode: bool = False
-     """Whether to run real-time tests."""
- 
-diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py
-index 3ea405770..5a42001b9 100644
---- a/megatron/core/transformer/transformer_layer.py
-+++ b/megatron/core/transformer/transformer_layer.py
-@@ -223,6 +223,7 @@ class TransformerLayerSubmodules:
-     input_layernorm: Union[ModuleSpec, type] = IdentityOp
-     self_attention: Union[ModuleSpec, type] = IdentityOp
-     self_attn_bda: Union[ModuleSpec, type] = IdentityFuncOp
-+    post_self_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
- 
-     pre_cross_attn_layernorm: Union[ModuleSpec, type] = IdentityOp
-     cross_attention: Union[ModuleSpec, type] = IdentityOp
-@@ -231,6 +232,7 @@ class TransformerLayerSubmodules:
-     pre_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
-     mlp: Union[ModuleSpec, type] = IdentityOp
-     mlp_bda: Union[ModuleSpec, type] = IdentityFuncOp
-+    post_mlp_layernorm: Union[ModuleSpec, type] = IdentityOp
- 
-     # Mapping for sharded tensor keys to be applied in `sharded_state_dict` method
-     sharded_state_dict_keys_map: Dict[str, str] = field(default_factory=dict)
-@@ -310,6 +312,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
-         # [Module 3: BiasDropoutFusion]
-         self.self_attn_bda = build_module(submodules.self_attn_bda)
- 
-+        self.post_self_attn_layernorm = build_module(
-+            submodules.post_self_attn_layernorm,
-+            config=self.config,
-+            hidden_size=self.config.hidden_size,
-+            eps=self.config.layernorm_epsilon,
-+        )
-+
-         # [Module 4: Post SelfAttention] Optional Layernorm after self-attn
-         self.pre_cross_attn_layernorm = build_module(
-             submodules.pre_cross_attn_layernorm,
-@@ -375,6 +384,13 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
- 
-         self.is_moe_layer = isinstance(self.mlp, MoELayer)
- 
-+        self.post_mlp_layernorm = build_module(
-+            submodules.post_mlp_layernorm,
-+            config=self.config,
-+            hidden_size=self.config.hidden_size,
-+            eps=self.config.layernorm_epsilon
-+        )
-+
-         self.recompute_input_layernorm = False
-         self.recompute_pre_mlp_layernorm = False
-         self.recompute_mlp = False
-@@ -551,6 +567,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
-                 attention_output_with_bias[0]
-             )
- 
-+        attention_output, attention_output_bias = attention_output_with_bias
-+        attention_output = self.post_self_attn_layernorm(attention_output)
-+        attention_output_with_bias = (attention_output, attention_output_bias)
-+
-         # TODO: could we move `bias_dropout_add_exec_handler` itself
-         # inside the module provided in the `bias_dropout_add_spec` module?
-         nvtx_range_push(suffix="self_attn_bda")
-@@ -677,6 +697,10 @@ class TransformerLayer(GraphableMegatronModule, BaseTransformerLayer):
-         else:
-             mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)
- 
-+        mlp_output, mlp_output_bias = mlp_output_with_bias
-+        mlp_output = self.post_mlp_layernorm(mlp_output)
-+        mlp_output_with_bias = (mlp_output, mlp_output_bias)
-+
-         if self.recompute_pre_mlp_layernorm:
-             # discard the output of the pre-mlp layernorm and register the recompute
-             # as a gradient hook of mlp_output_with_bias[0]
-diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
-index b267c8a81..83736acdc 100644
---- a/megatron/training/arguments.py
-+++ b/megatron/training/arguments.py
-@@ -1398,6 +1398,9 @@ def core_transformer_config_from_args(args, config_class=None):
- 
-     kw_args['inference_sampling_seed'] = args.seed
- 
-+    kw_args['post_self_attn_layernorm'] = args.post_self_attn_layernorm
-+    kw_args['post_mlp_layernorm'] = args.post_mlp_layernorm
-+
-     # handle quantization config
-     # NOTE: Kitchen arguments are only added to the namespace when
-     # Kitchen library is available.
-@@ -1764,6 +1767,12 @@ def _add_network_size_args(parser):
-                        action='store_true',
-                        help='If set, use original BERT residula connection '
-                        'ordering.')
-+    group.add_argument('--post-self-attn-layernorm', action='store_true',
-+                       help='If set, use post self attention layernorm.')
-+    group.add_argument('--post-mlp-layernorm', action='store_true',
-+                       help='If set, use post MLP layernorm.')
-+    group.add_argument('--use-gated-attention', action='store_true',
-+                       help='If set, use gated attention as in Qwen3Next')
-     group.add_argument('--openai-gelu', action='store_true',
-                        help='Use OpenAIs GeLU implementation. This option'
-                        'should not be used unless for backward compatibility'
-diff --git a/megatron/training/tokenizer/tokenizer.py b/megatron/training/tokenizer/tokenizer.py
-index 13b7526ca..6c590f653 100644
---- a/megatron/training/tokenizer/tokenizer.py
-+++ b/megatron/training/tokenizer/tokenizer.py
-@@ -136,7 +136,7 @@ class _HuggingFaceTokenizer(MegatronLegacyTokenizer):
-         # TODO(bnorick): download tokenizer once to lustre and use force offline to make sure all tasks read it from there
-         self._tokenizer = transformers.AutoTokenizer.from_pretrained(
-             pretrained_model_name_or_path=pretrained_model_name_or_path,
--            trust_remote_code=trust_remote_code,
-+            trust_remote_code=True,
-             **kwargs,
-         )
-         self._vocab = self._tokenizer.get_vocab()
diff --git a/docker/patch/v0.5.9/sglang.patch b/docker/patch/v0.5.9/sglang.patch
deleted file mode 100644
index e9145702e1..0000000000
--- a/docker/patch/v0.5.9/sglang.patch
+++ /dev/null
@@ -1,3216 +0,0 @@
-diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py
-index 6fbd1db823..f80ec11bb4 100644
---- a/python/sglang/srt/configs/model_config.py
-+++ b/python/sglang/srt/configs/model_config.py
-@@ -274,6 +274,7 @@ class ModelConfig:
- 
-         if is_draft_model and self.hf_config.architectures[0] in [
-             "DeepseekV3ForCausalLM",
-+            "DeepseekV32ForCausalLM",
-             "GlmMoeDsaForCausalLM",
-         ]:
-             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
---- a/python/sglang/srt/disaggregation/base/conn.py
-+++ b/python/sglang/srt/disaggregation/base/conn.py
-@@ -17,6 +17,7 @@ class KVArgs:
-     kv_data_ptrs: List[int]
-     kv_data_lens: List[int]
-     kv_item_lens: List[int]
-+    aux_buffer_names: List[str]
-     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
---- 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
- from __future__ import annotations
- 
- import logging
-+import os
- import time
- from collections import deque
- from dataclasses import dataclass
-@@ -40,8 +41,10 @@ from sglang.srt.disaggregation.utils import (
-     MetadataBuffers,
-     ReqToMetadataIdxAllocator,
-     TransferBackend,
-+    apply_prefill_timing_payload,
-     get_kv_class,
-     is_mla_backend,
-+    is_slime_profiling_enabled,
-     kv_to_page_indices,
-     poll_and_all_reduce,
-     prepare_abort,
-@@ -295,6 +298,7 @@ class DecodePreallocQueue:
-         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
-             self.metadata_buffers.get_buf_infos()
-         )
-+        kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names()
- 
-         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:
-         )
-         return kv_manager
- 
-+    def release_memory_occupation(self):
-+        self.queue.clear()
-+        self.retracted_queue.clear()
-+        if hasattr(self.kv_manager, "deregister_buffer_to_engine"):
-+            self.kv_manager.deregister_buffer_to_engine()
-+
-+    def resume_memory_occupation(self):
-+        if hasattr(self.kv_manager, "register_buffer_to_engine"):
-+            self.kv_manager.register_buffer_to_engine()
-+
-     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:
-             [decode_req.kv_receiver for decode_req in self.queue], self.gloo_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")
-+        )
-+        now = time.perf_counter()
-+
-         for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
-             if rids_to_check is not None and decode_req.req.rid not in rids_to_check:
-                 continue
- 
-             if poll == KVPoll.Bootstrapping:
--                pass
-+                # Check for bootstrap timeout
-+                entry_time = getattr(
-+                    decode_req.req.time_stats,
-+                    "decode_prealloc_queue_entry_time",
-+                    None,
-+                )
-+                if entry_time is not None and (now - entry_time) > bootstrap_timeout:
-+                    error_message = (
-+                        f"Decode bootstrap timed out after {now - entry_time:.1f}s "
-+                        f"for request rank={self.tp_rank} "
-+                        f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
-+                    )
-+                    logger.error(error_message)
-+                    prepare_abort(
-+                        decode_req.req,
-+                        error_message,
-+                        status_code=HTTPStatus.GATEWAY_TIMEOUT,
-+                    )
-+                    if self.scheduler.enable_metrics:
-+                        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:
-                 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(
-                 page_indices, decode_req.metadata_buffer_index, state_indices
-@@ -751,6 +791,7 @@ class DecodeTransferQueue:
-             output_topk_index,
-             output_hidden_states,
-             output_bootstrap_room,
-+            output_prefill_timing,
-         ) = self.metadata_buffers.get_buf(idx)
- 
-         # Validate bootstrap_room to detect context corruption
-@@ -813,6 +854,14 @@ class DecodeTransferQueue:
-                 output_top_logprobs_idx[: decode_req.req.top_logprobs_num].tolist()
-             )
- 
-+        # Inject prefill-side PD timing forwarded from the P instance.
-+        # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap,
-+        #          alloc_waiting, transfer_speed, transfer_mb, retry_count]
-+        if is_slime_profiling_enabled():
-+            apply_prefill_timing_payload(
-+                decode_req.req.time_stats, output_prefill_timing
-+            )
-+
-         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
-         )
- 
-+        # 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.
-+        transfer_timeout = float(
-+            os.environ.get("SGLANG_DISAGGREGATION_TRANSFER_TIMEOUT", "600")
-+        )
-+        now = time.perf_counter()
-+
-         transferred_reqs = []
-         indices_to_remove = set()
-         for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
-@@ -877,7 +933,20 @@ class DecodeTransferQueue:
-                 KVPoll.WaitingForInput,
-                 KVPoll.Transferring,
-             ]:
--                pass
-+                # Check for transfer timeout
-+                entry_time = getattr(
-+                    decode_req.req.time_stats,
-+                    "decode_transfer_queue_entry_time",
-+                    None,
-+                )
-+                if entry_time is not None and (now - entry_time) > transfer_timeout:
-+                    error_message = (
-+                        f"Decode transfer timed out after {now - entry_time:.1f}s "
-+                        f"(state={poll}) for request rank={self.tp_rank} "
-+                        f"{decode_req.req.rid=} {decode_req.req.bootstrap_room=}"
-+                    )
-+                    logger.error(error_message)
-+                    decode_req.kv_receiver.abort()
-             else:
-                 raise ValueError(f"Unexpected poll case: {poll}")
- 
-@@ -893,6 +962,14 @@ class DecodeTransferQueue:
- 
-         return transferred_reqs
- 
-+    def release_memory_occupation(self):
-+        """Clean up all in-flight transfers before releasing GPU memory."""
-+        self.queue.clear()
-+
-+    def resume_memory_occupation(self):
-+        """Resume after GPU memory re-allocation. Queue was already cleared on release."""
-+        pass
-+
- 
- class SchedulerDisaggregationDecodeMixin:
- 
-@@ -1072,7 +1149,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:
--            # if there are still retracted requests, we do not allocate new requests
-+            # Still have retracted requests that couldn't resume (not enough memory).
-+            # Don't accept new requests (pop_preallocated) — they would consume memory
-+            # that retracted requests need.
-+            # But DO drain completed transfers: their KV is already committed, and
-+            # moving them to waiting_queue frees the reserved-decode-token budget
-+            # in _allocatable_tokens(), which may unblock resume on the next iteration.
-+            # Without this, completed transfers hold memory indefinitely → deadlock.
-+            alloc_reqs = self.disagg_decode_transfer_queue.pop_transferred()
-+            self.waiting_queue.extend(alloc_reqs)
-             return
- 
-         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
---- 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,
- )
--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):
-                 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):
-         prefill_aux_ptrs = self.kv_args.aux_data_ptrs
-         prefill_aux_item_lens = self.kv_args.aux_item_lens
- 
--        for i, dst_aux_ptr in enumerate(dst_aux_ptrs):
--            length = prefill_aux_item_lens[i]
--            src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
--            dst_addr = dst_aux_ptrs[i] + length * req.dst_aux_index
-+        for _, src_addr, dst_addr, length in iter_aux_transfer_specs(
-+            self.kv_args.aux_buffer_names,
-+            prefill_aux_ptrs,
-+            prefill_aux_item_lens,
-+            dst_aux_ptrs,
-+            prefill_aux_index,
-+            req.dst_aux_index,
-+        ):
-             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):
-         prefill_aux_ptrs = self.kv_args.aux_data_ptrs
-         prefill_aux_item_lens = self.kv_args.aux_item_lens
- 
--        for i in range(len(prefill_aux_ptrs)):
--            length = prefill_aux_item_lens[i]
--            src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
-+        for i, src_addr, _, length in iter_aux_transfer_specs(
-+            self.kv_args.aux_buffer_names,
-+            prefill_aux_ptrs,
-+            prefill_aux_item_lens,
-+            dst_aux_ptrs,
-+            prefill_aux_index,
-+            req.dst_aux_index,
-+        ):
-             data = AuxDataCodec.serialize_data_from_buffer(src_addr, length)
- 
-             self.send_aux_data_to_endpoint(
-@@ -643,13 +665,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."
-                 )
--            if len(prefill_state_indices) < len(req.dst_state_indices):
--                logger.warning(
--                    f"len(prefill_state_indices) = {len(prefill_state_indices)}, len(dst_state_indices) = {len(req.dst_state_indices)}"
-+            if len(prefill_state_indices) != len(req.dst_state_indices):
-+                logger.error(
-+                    "PD extra-state index mismatch, reject transfer to avoid corrupted outputs: "
-+                    f"len(prefill_state_indices)={len(prefill_state_indices)}, "
-+                    f"len(dst_state_indices)={len(req.dst_state_indices)}"
-                 )
--                prefill_state_indices = prefill_state_indices[
--                    : len(req.dst_state_indices)
--                ]
-+                return -1
-             # 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):
-                         if ret != 0:
-                             with self.session_lock:
-                                 self.session_failures[req.mooncake_session_id] += 1
--                                # Failures should never happen if the session is not dead, if the session fails once, mark it as failed
--                                if self.session_failures[req.mooncake_session_id] >= 1:
--                                    self.failed_sessions.add(req.mooncake_session_id)
--                                    logger.error(
--                                        f"Session {req.mooncake_session_id} failed."
--                                    )
-                             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):
- 
-                         if kv_chunk.is_last:
-                             if kv_chunk.state_indices is not None:
--                                self.maybe_send_extra(
-+                                ret = self.maybe_send_extra(
-                                     req,
-                                     kv_chunk.state_indices,
-                                     target_rank_registration_info.dst_state_data_ptrs,
-                                     executor,
-                                     target_rank_registration_info,
-                                 )
-+                                if ret != 0:
-+                                    with self.session_lock:
-+                                        self.session_failures[
-+                                            req.mooncake_session_id
-+                                        ] += 1
-+                                    self.record_failure(
-+                                        kv_chunk.room,
-+                                        f"Failed to send extra state chunk of {kv_chunk.room} to {req.endpoint}:{req.dst_port}",
-+                                    )
-+                                    self.update_status(kv_chunk.room, KVPoll.Failed)
-+                                    self.sync_status_to_decode_endpoint(
-+                                        req.endpoint,
-+                                        req.dst_port,
-+                                        req.room,
-+                                        KVPoll.Failed,
-+                                        local_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
---- 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
- from __future__ import annotations
- 
- import logging
-+import os
- import time
- from collections import deque
- from http import HTTPStatus
-@@ -167,6 +168,7 @@ class PrefillBootstrapQueue:
-         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
-             self.metadata_buffers.get_buf_infos()
-         )
-+        kv_args.aux_buffer_names = self.metadata_buffers.get_aux_buffer_names()
-         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
-         )
- 
-+        # 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")
-+        )
-+        now = time.perf_counter()
-+
-         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:
-                     continue
- 
-             if poll == KVPoll.Bootstrapping:
-+                # Check for bootstrap timeout
-+                entry_time = getattr(
-+                    req.time_stats,
-+                    "prefill_bootstrap_queue_entry_time",
-+                    None,
-+                )
-+                if entry_time is not None and (now - entry_time) > bootstrap_timeout:
-+                    error_message = (
-+                        f"Prefill bootstrap timed out after {now - entry_time:.1f}s "
-+                        f"for request rank={self.tp_rank} "
-+                        f"{req.rid=} {req.bootstrap_room=}"
-+                    )
-+                    logger.error(error_message)
-+                    prepare_abort(
-+                        req, error_message, status_code=HTTPStatus.GATEWAY_TIMEOUT
-+                    )
-+                    self.scheduler.stream_output([req], req.return_logprob)
-+                    indices_to_remove.add(i)
-+                    failed_reqs.append(req)
-+                    if self.scheduler.enable_metrics:
-+                        self.scheduler.metrics_collector.increment_bootstrap_failed_reqs()
-                 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:
-         else:
-             return bootstrapped_reqs, failed_reqs
- 
-+    def release_memory_occupation(self):
-+        self.queue.clear()
-+        if hasattr(self.kv_manager, "deregister_buffer_to_engine"):
-+            self.kv_manager.deregister_buffer_to_engine()
-+
-+    def resume_memory_occupation(self):
-+        if hasattr(self.kv_manager, "register_buffer_to_engine"):
-+            self.kv_manager.register_buffer_to_engine()
-+
- 
- class SchedulerDisaggregationPrefillMixin:
-     """
-@@ -547,6 +585,18 @@ class SchedulerDisaggregationPrefillMixin:
- 
-         self.maybe_send_health_check_signal()
- 
-+        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(
-+                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:
-             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")
-+        )
-+        now = time.perf_counter()
-+
-         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)
-                     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",
-+                    None,
-+                )
-+                if entry_time is not None and (now - entry_time) > transfer_timeout:
-+                    error_message = (
-+                        f"Prefill transfer timed out after {now - entry_time:.1f}s "
-+                        f"(state={poll}) for request rank={self.tp_rank} "
-+                        f"{req.rid=} {req.bootstrap_room=}"
-+                    )
-+                    logger.error(error_message)
-+                    release_kv_cache(req, self.tree_cache)  # unlock the tree
-+                    prepare_abort(
-+                        req, error_message, status_code=HTTPStatus.GATEWAY_TIMEOUT
-+                    )
-+                    if hasattr(req.disagg_kv_sender, "clear"):
-+                        req.disagg_kv_sender.clear()
-+                    done_reqs.append(req)
-+                    if self.enable_metrics:
-+                        self.metrics_collector.increment_transfer_failed_reqs()
-+                else:
-+                    undone_reqs.append(req)
-             elif poll == KVPoll.Success:  # transfer done
-                 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
---- a/python/sglang/srt/disaggregation/utils.py
-+++ b/python/sglang/srt/disaggregation/utils.py
-@@ -21,6 +21,17 @@ if TYPE_CHECKING:
- # Constants & Enums
- #########################
- FAKE_BOOTSTRAP_HOST = "2.2.2.2"
-+PREFILL_TIMING_AUX_BUFFER_NAME = "prefill_timing"
-+PREFILL_TIMING_DEST_ATTRS = (
-+    ("fwd_prefill_bootstrap_queue_duration", float),
-+    ("fwd_prefill_forward_duration", float),
-+    ("fwd_prefill_transfer_queue_duration", float),
-+    ("fwd_bootstrap_duration", float),
-+    ("fwd_alloc_waiting_duration", float),
-+    ("fwd_transfer_speed_gb_s", float),
-+    ("fwd_transfer_total_mb", float),
-+    ("fwd_prefill_retry_count", int),
-+)
- 
- 
- class DisaggregationMode(Enum):
-@@ -139,46 +150,35 @@ class MetadataBuffers:
-             self.bootstrap_room = torch.zeros(
-                 (size, 8), dtype=torch.uint64, device=device
-             )
-+            # Prefill-side PD timing (8 floats, padded to 16 for RDMA alignment).
-+            # Layout: [bootstrap_queue, forward, transfer_queue, bootstrap,
-+            #          alloc_waiting, transfer_speed, transfer_mb, retry_count]
-+            self.prefill_timing = torch.zeros(
-+                (size, 16), dtype=torch.float32, device=device
-+            )
-+        self.aux_buffers = [
-+            ("output_ids", self.output_ids),
-+            ("cached_tokens", self.cached_tokens),
-+            ("output_token_logprobs_val", self.output_token_logprobs_val),
-+            ("output_token_logprobs_idx", self.output_token_logprobs_idx),
-+            ("output_top_logprobs_val", self.output_top_logprobs_val),
-+            ("output_top_logprobs_idx", self.output_top_logprobs_idx),
-+            ("output_topk_p", self.output_topk_p),
-+            ("output_topk_index", self.output_topk_index),
-+            ("output_hidden_states", self.output_hidden_states),
-+            ("bootstrap_room", self.bootstrap_room),
-+            (PREFILL_TIMING_AUX_BUFFER_NAME, self.prefill_timing),
-+        ]
- 
-     def get_buf_infos(self):
--        ptrs = [
--            self.output_ids.data_ptr(),
--            self.cached_tokens.data_ptr(),
--            self.output_token_logprobs_val.data_ptr(),
--            self.output_token_logprobs_idx.data_ptr(),
--            self.output_top_logprobs_val.data_ptr(),
--            self.output_top_logprobs_idx.data_ptr(),
--            self.output_topk_p.data_ptr(),
--            self.output_topk_index.data_ptr(),
--            self.output_hidden_states.data_ptr(),
--            self.bootstrap_room.data_ptr(),
--        ]
--        data_lens = [
--            self.output_ids.nbytes,
--            self.cached_tokens.nbytes,
--            self.output_token_logprobs_val.nbytes,
--            self.output_token_logprobs_idx.nbytes,
--            self.output_top_logprobs_val.nbytes,
--            self.output_top_logprobs_idx.nbytes,
--            self.output_topk_p.nbytes,
--            self.output_topk_index.nbytes,
--            self.output_hidden_states.nbytes,
--            self.bootstrap_room.nbytes,
--        ]
--        item_lens = [
--            self.output_ids[0].nbytes,
--            self.cached_tokens[0].nbytes,
--            self.output_token_logprobs_val[0].nbytes,
--            self.output_token_logprobs_idx[0].nbytes,
--            self.output_top_logprobs_val[0].nbytes,
--            self.output_top_logprobs_idx[0].nbytes,
--            self.output_topk_p[0].nbytes,
--            self.output_topk_index[0].nbytes,
--            self.output_hidden_states[0].nbytes,
--            self.bootstrap_room[0].nbytes,
--        ]
-+        ptrs = [buffer.data_ptr() for _, buffer in self.aux_buffers]
-+        data_lens = [buffer.nbytes for _, buffer in self.aux_buffers]
-+        item_lens = [buffer[0].nbytes for _, buffer in self.aux_buffers]
-         return ptrs, data_lens, item_lens
- 
-+    def get_aux_buffer_names(self):
-+        return [name for name, _ in self.aux_buffers]
-+
-     def get_buf(self, idx: int):
-         return (
-             self.output_ids[idx],
-@@ -191,8 +191,12 @@ class MetadataBuffers:
-             self.output_topk_index[idx],
-             self.output_hidden_states[idx],
-             self.bootstrap_room[idx],
-+            self.prefill_timing[idx],
-         )
- 
-+    def clear_profiling_buf(self, idx: int):
-+        self.prefill_timing[idx].zero_()
-+
-     def set_buf(self, req: Req):
- 
-         self.output_ids[req.metadata_buffer_index][0] = req.output_ids[0]
-@@ -237,6 +241,84 @@ class MetadataBuffers:
-         self.bootstrap_room[req.metadata_buffer_index, 0] = (
-             req.bootstrap_room if req.bootstrap_room is not None else 0
-         )
-+        # Pack prefill-side PD timing durations for transfer to decode instance.
-+        # Note: set_buf is called at the START of the last KV chunk send, so
-+        # completion_time and prefill_transfer_queue_entry_time are not yet set.
-+        # We use time.perf_counter() as the "forward just completed" timestamp.
-+        import time
-+
-+        ts = req.time_stats
-+        timing = self.prefill_timing[req.metadata_buffer_index]
-+        self.clear_profiling_buf(req.metadata_buffer_index)
-+        if not is_slime_profiling_enabled():
-+            return
-+        for idx, value in enumerate(
-+            build_prefill_timing_payload(ts, now=time.perf_counter())
-+        ):
-+            if value > 0:
-+                timing[idx] = value
-+
-+
-+def is_slime_profiling_enabled() -> bool:
-+    return envs.SLIME_ENABLE_PROFILING.get()
-+
-+
-+def build_prefill_timing_payload(time_stats, now: float) -> tuple[float, ...]:
-+    bootstrap_queue_duration = 0.0
-+    if (
-+        time_stats.prefill_bootstrap_queue_entry_time > 0
-+        and time_stats.wait_queue_entry_time > 0
-+    ):
-+        bootstrap_queue_duration = (
-+            time_stats.wait_queue_entry_time
-+            - time_stats.prefill_bootstrap_queue_entry_time
-+        )
-+
-+    prefill_forward_duration = (
-+        now - time_stats.forward_entry_time
-+        if time_stats.forward_entry_time > 0
-+        else 0.0
-+    )
-+
-+    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, time_stats.transfer_speed_gb_s),
-+        max(0.0, time_stats.transfer_total_mb),
-+        float(max(0, time_stats.prefill_retry_count)),
-+    )
-+
-+
-+def apply_prefill_timing_payload(time_stats, timing) -> None:
-+    for value, (attr_name, caster) in zip(
-+        timing[: len(PREFILL_TIMING_DEST_ATTRS)].tolist(),
-+        PREFILL_TIMING_DEST_ATTRS,
-+    ):
-+        if value > 0:
-+            setattr(time_stats, attr_name, caster(value))
-+
-+
-+def iter_aux_transfer_specs(
-+    aux_buffer_names: list[str],
-+    prefill_aux_ptrs: list[int],
-+    prefill_aux_item_lens: list[int],
-+    dst_aux_ptrs: list[int],
-+    prefill_aux_index: int,
-+    dst_aux_index: int,
-+):
-+    profiling_enabled = is_slime_profiling_enabled()
-+    for i, (buffer_name, dst_aux_ptr) in enumerate(zip(aux_buffer_names, dst_aux_ptrs)):
-+        if not profiling_enabled and buffer_name == PREFILL_TIMING_AUX_BUFFER_NAME:
-+            continue
-+        length = prefill_aux_item_lens[i]
-+        if length <= 0:
-+            continue
-+        src_addr = prefill_aux_ptrs[i] + length * prefill_aux_index
-+        dst_addr = dst_aux_ptr + length * dst_aux_index
-+        yield i, src_addr, dst_addr, length
- 
- 
- #########################
-diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
-index 0ed5a1b44b..67e33c650d 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 (
-     LoadLoRAAdapterReqInput,
-     MultimodalDataInputFormat,
-     OpenSessionReqInput,
-+    PostProcessWeightsReqInput,
-     ReleaseMemoryOccupationReqInput,
-     ResumeMemoryOccupationReqInput,
-     RpcReqInput,
-@@ -641,6 +642,24 @@ class Engine(EngineBase):
-             self.tokenizer_manager.update_weights_from_ipc(obj, None)
-         )
- 
-+    def post_process_weights(
-+        self,
-+        restore_weights_before_load: bool = False,
-+        post_process_quantization: bool = False,
-+    ):
-+        """
-+        Optional post-processing for updated weights (e.g., Marlin conversion).
-+        Should be called after weight update is finished.
-+        """
-+        obj = PostProcessWeightsReqInput(
-+            restore_weights_before_load=restore_weights_before_load,
-+            post_process_quantization=post_process_quantization,
-+        )
-+
-+        return self.loop.run_until_complete(
-+            self.tokenizer_manager.post_process_weights(obj, None)
-+        )
-+
-     def get_weights_by_name(self, name: str, truncate_size: int = 100):
-         """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
---- 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 (
-     OpenSessionReqInput,
-     ParseFunctionCallReq,
-     PauseGenerationReqInput,
-+    PostProcessWeightsReqInput,
-     ProfileReqInput,
-     ReleaseMemoryOccupationReqInput,
-     ResumeMemoryOccupationReqInput,
-@@ -574,10 +575,8 @@ async def model_info():
- @app.get("/weight_version")
- async def weight_version():
-     """Get the current weight version."""
--    raise HTTPException(
--        status_code=404,
--        detail="Endpoint '/get_weight_version' or '/weight_version' is deprecated. Please use '/model_info' instead.",
--    )
-+    result = await model_info()
-+    return {"weight_version": result.get("weight_version", None)}
- 
- 
- @app.get("/get_server_info")
-@@ -594,9 +593,19 @@ async def get_server_info():
- async def server_info():
-     """Get the server information."""
-     # Returns internal states per DP.
--    internal_states: List[Dict[Any, Any]] = (
--        await _global_state.tokenizer_manager.get_internal_state()
--    )
-+    # In large/disaggregated deployments this can occasionally block; keep endpoint responsive.
-+    server_info_timeout = float(os.environ.get("SGLANG_SERVER_INFO_TIMEOUT", "2"))
-+    try:
-+        internal_states: List[Dict[Any, Any]] = await asyncio.wait_for(
-+            _global_state.tokenizer_manager.get_internal_state(),
-+            timeout=server_info_timeout,
-+        )
-+    except asyncio.TimeoutError:
-+        logger.warning(
-+            "Timed out getting internal state for /server_info after %.1fs; returning empty internal_states",
-+            server_info_timeout,
-+        )
-+        internal_states = []
- 
-     # 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
-         return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
- 
- 
-+@app.post("/post_process_weights")
-+@auth_level(AuthLevel.ADMIN_OPTIONAL)
-+async def post_process_weights(req: PostProcessWeightsReqInput, request: Request):
-+    """
-+    Optional post-processing for updated weights (e.g., Marlin conversion).
-+    This should be called selectively after `update_weights_from_distributed/update_weights_from_tensor`.
-+    """
-+    success, message = await _global_state.tokenizer_manager.post_process_weights(
-+        req, request
-+    )
-+
-+    content = {"success": success, "message": message}
-+    return ORJSONResponse(
-+        content, status_code=200 if success else HTTPStatus.BAD_REQUEST
-+    )
-+
-+
- @app.post("/update_weight_version")
- @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
---- a/python/sglang/srt/environ.py
-+++ b/python/sglang/srt/environ.py
-@@ -244,6 +244,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,
-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
---- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
-+++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py
-@@ -1,6 +1,7 @@
- from __future__ import annotations
- 
- import contextlib
-+import os
- from abc import ABC, abstractmethod
- from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
- 
-@@ -201,14 +202,31 @@ class Indexer(MultiPlatformOp):
-             prefix=add_prefix("weights_proj", prefix),
-         )
-         self.k_norm = LayerNorm(self.head_dim, dtype=torch.float32)
-+        server_args = get_global_server_args()
-+        disable_flag = server_args.disable_indexer_rope_neox_style
-+        env_raw = os.environ.get("INDEXER_ROPE_NEOX_STYLE", None)
-+        if env_raw is not None:
-+            env_value = env_raw == "1"
-+            if disable_flag and env_value:
-+                raise ValueError(
-+                    "Conflict: --disable-indexer-rope-neox-style is set but "
-+                    "INDEXER_ROPE_NEOX_STYLE='1'. "
-+                    "Please remove one or make them consistent."
-+                )
-+            resolved_neox_style = env_value
-+        elif disable_flag:
-+            resolved_neox_style = False
-+        else:
-+            resolved_neox_style = is_neox_style
-+
-         self.rotary_emb = get_rope_wrapper(
-             rope_head_dim,
-             rotary_dim=rope_head_dim,
-             max_position=max_position_embeddings,
-             base=rope_theta,  # type: ignore
-             rope_scaling=rope_scaling,
--            is_neox_style=is_neox_style,
--            device=get_global_server_args().device,
-+            is_neox_style=resolved_neox_style,
-+            device=server_args.device,
-         )
-         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()
-+        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 * self.n_heads**-0.5
-         weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
-         return weights
-@@ -982,15 +1005,26 @@ class Indexer(MultiPlatformOp):
-             query, key = self._get_q_k_bf16(
-                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
-             )
-+            if query.shape[1] < 32:
-+                assert 32 % query.shape[1] == 0
-+                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:
-             query, key = self._get_q_k_bf16(
-                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
-             )
-+            if query.shape[1] < 32:
-+                assert 32 % query.shape[1] == 0
-+                query = query.repeat_interleave(32 // query.shape[1], dim=1)
- 
-             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
---- 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):
-                     "CompressedTensorsWNA16TritonMoE",
-                 ]
-             )
-+            and "zero" not in weight_name
-             else loaded_weight
-         )
- 
-@@ -916,6 +917,7 @@ class FusedMoE(torch.nn.Module):
-                     "CompressedTensorsWNA16TritonMoE",
-                 ]
-             )
-+            and "zero" not in weight_name
-             else loaded_weight
-         )
- 
-diff --git a/python/sglang/srt/layers/moe/routed_experts_capturer.py b/python/sglang/srt/layers/moe/routed_experts_capturer.py
-index 00bd687555..12d5577af2 100644
---- a/python/sglang/srt/layers/moe/routed_experts_capturer.py
-+++ b/python/sglang/srt/layers/moe/routed_experts_capturer.py
-@@ -8,10 +8,15 @@ import torch
- 
- from sglang.srt.configs.model_config import ModelConfig
- from sglang.srt.layers.dp_attention import (
-+    attn_tp_all_gather_into_tensor,
-     get_attention_dp_rank,
-+    get_attention_tp_size,
-     get_dp_local_info,
-     is_dp_attention_enabled,
- )
-+from sglang.srt.layers.moe import (
-+    get_moe_a2a_backend,
-+)
- from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
- from sglang.srt.model_executor.forward_batch_info import ForwardBatch
- from sglang.srt.server_args import get_global_server_args
-@@ -181,13 +186,26 @@ class _RoutedExpertsCapturerReal(RoutedExpertsCapturer):
-             device=device,
-         )
- 
-+        if get_moe_a2a_backend().is_deepep():
-+            attn_tp_size = get_attention_tp_size() if is_dp_attention_enabled() else 1
-+            self.gather_buffer = torch.empty(
-+                (
-+                    self.device_cache.buffer.shape[0] * attn_tp_size,
-+                    self.device_cache.buffer.shape[2],
-+                ),
-+                dtype=torch.int32,
-+                device=device,
-+            )
-+
-     def _sync_fwd_experts_buffer_DtoH(
-         self,
-         forward_batch: ForwardBatch,
-         can_run_graph: bool,
-         cuda_graph_batch: int,
-     ):
--        if is_dp_attention_enabled():
-+        # When DeepEP is enabled, capture() already does all_gather, so device_cache.buffer
-+        # contains data from all DP ranks. We should not slice by DP rank in this case.
-+        if is_dp_attention_enabled() and not get_moe_a2a_backend().is_deepep():
-             local_start_pos, local_num_tokens = get_dp_local_info(forward_batch)
-             # handle with cuda graph padding
-             if can_run_graph:
-@@ -206,6 +224,12 @@ class _RoutedExpertsCapturerReal(RoutedExpertsCapturer):
-         ].cpu()
- 
-     def capture(self, layer_id: int, topk_ids: torch.Tensor):
-+        if get_moe_a2a_backend().is_deepep():
-+            local_topk_ids = topk_ids
-+            topk_ids = self.gather_buffer[
-+                : local_topk_ids.size(0) * get_attention_tp_size()
-+            ]
-+            attn_tp_all_gather_into_tensor(topk_ids, local_topk_ids)
-         self.device_cache.capture_fwd_routed_experts(layer_id, topk_ids)
- 
-     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
---- 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):
-         )
-         is_static = not weight_quant.dynamic
- 
--        return is_channel_group and input_quant_none and is_symmetric and is_static
-+        return is_channel_group and input_quant_none and is_static
- 
-     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):
-     def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
-         layer.scheme.process_weights_after_loading(layer)
- 
-+    def restore_weights_before_loading(self, layer: torch.nn.Module) -> None:
-+        layer.scheme.restore_weights_before_loading(layer)
-+
-     def create_weights(
-         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
---- 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 (
-     CompressedTensorsMoEScheme,
- )
- from sglang.srt.layers.quantization.gptq import gptq_marlin_moe_repack
--from sglang.srt.layers.quantization.marlin_utils import marlin_moe_permute_scales
-+from sglang.srt.layers.quantization.marlin_utils import (
-+    marlin_moe_permute_scales,
-+    moe_awq_to_marlin_zero_points,
-+)
- from sglang.srt.layers.quantization.utils import replace_parameter
- from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, set_weight_attrs
- 
-@@ -64,7 +67,7 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
-         self.strategy = config.strategy
-         self.group_size = config.group_size
-         self.actorder = config.actorder
--        assert config.symmetric, "Only symmetric quantization is supported for MoE"
-+        self.sym = config.symmetric
- 
-         if not (
-             self.quant_config.quant_format == CompressionFormat.pack_quantized.value
-@@ -124,7 +127,7 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
- 
-         # In the case where we have actorder/g_idx,
-         # we do not partition the w2 scales
--        load_full_w2 = self.actorder and self.group_size != -1
-+        load_full_w2 = (self.actorder != "static") and self.group_size != -1
- 
-         if load_full_w2:
-             w2_scales_size = intermediate_size_per_partition * layer.moe_tp_size
-@@ -172,6 +175,32 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
-         layer.register_parameter("w13_weight_shape", w13_weight_shape)
-         set_weight_attrs(w13_weight_shape, extra_weight_attrs)
- 
-+        # add zero param
-+        if not self.sym:
-+            w13_qzeros = torch.nn.Parameter(
-+                torch.empty(
-+                    num_experts,
-+                    num_groups_w13,
-+                    2 * intermediate_size_per_partition // self.packed_factor,
-+                    dtype=torch.int32,
-+                ),
-+                requires_grad=False,
-+            )
-+            layer.register_parameter("w13_weight_zero_point", w13_qzeros)
-+            set_weight_attrs(w13_qzeros, extra_weight_attrs)
-+
-+            w2_qzeros = torch.nn.Parameter(
-+                torch.empty(
-+                    num_experts,
-+                    num_groups_w2,
-+                    hidden_size // self.packed_factor,
-+                    dtype=torch.int32,
-+                ),
-+                requires_grad=False,
-+            )
-+            layer.register_parameter("w2_weight_zero_point", w2_qzeros)
-+            set_weight_attrs(w2_qzeros, extra_weight_attrs)
-+
-         w13_g_idx = torch.nn.Parameter(
-             torch.empty(
-                 num_experts,
-@@ -225,11 +254,14 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
- 
-         # 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):
-         )
-         replace_tensor("w2_weight_scale", marlin_w2_scales)
- 
-+        # Repack zero
-+        if not self.sym:
-+            marlin_w13_zp = moe_awq_to_marlin_zero_points(
-+                layer.w13_weight_zero_point,
-+                size_k=layer.w13_weight_zero_point.shape[1],
-+                size_n=layer.w13_weight_zero_point.shape[2] * self.packed_factor,
-+                num_bits=self.num_bits,
-+            )
-+            replace_tensor("w13_weight_zero_point", marlin_w13_zp)
-+
-+            marlin_w2_zp = moe_awq_to_marlin_zero_points(
-+                layer.w2_weight_zero_point,
-+                size_k=layer.w2_weight_zero_point.shape[1],
-+                size_n=layer.w2_weight_zero_point.shape[2] * self.packed_factor,
-+                num_bits=self.num_bits,
-+            )
-+            replace_tensor("w2_weight_zero_point", marlin_w2_zp)
-+
-         layer.is_marlin_converted = True
- 
-     def restore_weights_before_loading(self, layer: torch.nn.Module):
-@@ -399,6 +449,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,
-+            w1_zeros=layer.w13_weight_zero_point if not self.sym else None,
-+            w2_zeros=layer.w2_weight_zero_point if not self.sym else None,
-             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
---- 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
- 
- 
-+@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
-+
-+
-+@dataclass
-+class PostProcessWeightsReqOutput(BaseReq):
-+    success: bool
-+    message: str
-+
-+
- @dataclass
- class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq):
-     success: bool
-@@ -1802,6 +1852,10 @@ 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
---- 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):
-             "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
-+
-+        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
-+            except (AttributeError, OSError):
-+                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,
-+            }
-+            return _uv_sub.spawn.Process(
-+                target=_uv_sub.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
-+    except Exception:
-+        pass
-+
- 
- 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
---- a/python/sglang/srt/managers/schedule_batch.py
-+++ b/python/sglang/srt/managers/schedule_batch.py
-@@ -1869,7 +1869,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
-         while first_iter or (
-             not self.check_decode_mem(selected_indices=sorted_indices)
-         ):
--            if len(sorted_indices) == 1:
-+            # We should allow all requests to be retracted in decode disaggregation mode
-+            # because there call be prealloc prefill requests.
-+            num_minimum_reqs = 0 if server_args.disaggregation_mode == "decode" else 1
-+            if len(sorted_indices) == num_minimum_reqs:
-                 # Always keep at least one request
-                 break
- 
-diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
-index a9ff0ac94b..a50dd5122b 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 (
-     OpenSessionReqInput,
-     OpenSessionReqOutput,
-     PauseGenerationReqInput,
-+    PostProcessWeightsReqInput,
-     ProfileReq,
-     ReleaseMemoryOccupationReqInput,
-     ResumeMemoryOccupationReqInput,
-@@ -1063,6 +1064,7 @@ class Scheduler(
-                 ),
-                 (UpdateWeightsFromTensorReqInput, self.update_weights_from_tensor),
-                 (UpdateWeightsFromIPCReqInput, self.update_weights_from_ipc),
-+                (PostProcessWeightsReqInput, self.post_process_weights),
-                 (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
---- 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()
- 
-         # 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,
-diff --git a/python/sglang/srt/managers/scheduler_profiler_mixin.py b/python/sglang/srt/managers/scheduler_profiler_mixin.py
-index 7d08f12b35..afc045da20 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:
-                 if self.profiler_prefill_ct > self.profiler_target_prefill_ct:
-                     if self.profile_in_progress:
-                         self.stop_profile(stage=ForwardMode.EXTEND)
--            elif batch.forward_mode.is_decode():
-+            elif batch.forward_mode.is_decode() or batch.forward_mode.is_prebuilt():
-                 if self.profiler_decode_ct == 0:
-                     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
---- 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 (
-     GPU_MEMORY_TYPE_KV_CACHE,
-     GPU_MEMORY_TYPE_WEIGHTS,
- )
-+from sglang.srt.disaggregation.utils import DisaggregationMode
- from sglang.srt.managers.io_struct import (
-     CheckWeightsReqInput,
-     CheckWeightsReqOutput,
-@@ -21,6 +22,8 @@ from sglang.srt.managers.io_struct import (
-     GetWeightsByNameReqOutput,
-     InitWeightsUpdateGroupReqInput,
-     InitWeightsUpdateGroupReqOutput,
-+    PostProcessWeightsReqInput,
-+    PostProcessWeightsReqOutput,
-     ReleaseMemoryOccupationReqInput,
-     ReleaseMemoryOccupationReqOutput,
-     ResumeMemoryOccupationReqInput,
-@@ -114,6 +117,11 @@ class SchedulerUpdateWeightsMixin:
-         torch.distributed.barrier(group=self.tp_cpu_group)
-         return UpdateWeightsFromIPCReqOutput(success, message)
- 
-+    def post_process_weights(self, recv_req: PostProcessWeightsReqInput):
-+        """Optional post-processing for updated weights (e.g., Marlin conversion)."""
-+        success, message = self.tp_worker.post_process_weights(recv_req)
-+        return PostProcessWeightsReqOutput(success, message)
-+
-     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:
-             self.memory_saver_adapter.pause(GPU_MEMORY_TYPE_KV_CACHE)
-             self.flush_cache()
- 
-+            if self.disaggregation_mode == DisaggregationMode.DECODE:
-+                if hasattr(self, "disagg_decode_transfer_queue"):
-+                    self.disagg_decode_transfer_queue.release_memory_occupation()
-+                if hasattr(self, "disagg_decode_prealloc_queue"):
-+                    self.disagg_decode_prealloc_queue.release_memory_occupation()
-+            elif self.disaggregation_mode == DisaggregationMode.PREFILL:
-+                if hasattr(self, "disagg_prefill_bootstrap_queue"):
-+                    self.disagg_prefill_bootstrap_queue.release_memory_occupation()
-+
-         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:
-         if GPU_MEMORY_TYPE_KV_CACHE in tags:
-             self.memory_saver_adapter.resume(GPU_MEMORY_TYPE_KV_CACHE)
- 
-+            if self.disaggregation_mode == DisaggregationMode.DECODE:
-+                if hasattr(self, "disagg_decode_transfer_queue"):
-+                    self.disagg_decode_transfer_queue.resume_memory_occupation()
-+                if hasattr(self, "disagg_decode_prealloc_queue"):
-+                    self.disagg_decode_prealloc_queue.resume_memory_occupation()
-+            elif self.disaggregation_mode == DisaggregationMode.PREFILL:
-+                if hasattr(self, "disagg_prefill_bootstrap_queue"):
-+                    self.disagg_prefill_bootstrap_queue.resume_memory_occupation()
-+
-         return ResumeMemoryOccupationReqOutput()
- 
-     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
---- 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 (
-     LoadLoRAAdapterReqOutput,
-     LoRAUpdateOutput,
-     OpenSessionReqInput,
-+    PostProcessWeightsReqInput,
-+    PostProcessWeightsReqOutput,
-     ProfileReq,
-     ProfileReqOutput,
-     ProfileReqType,
-@@ -187,6 +189,9 @@ class TokenizerCommunicatorMixin:
-         self.update_weights_from_ipc_communicator = _Communicator(
-             self.send_to_scheduler, server_args.dp_size
-         )
-+        self.post_process_weights_communicator = _Communicator(
-+            self.send_to_scheduler, server_args.dp_size
-+        )
-         self.get_weights_by_name_communicator = _Communicator(
-             self.send_to_scheduler, server_args.dp_size
-         )
-@@ -272,6 +277,10 @@ class TokenizerCommunicatorMixin:
-                     UpdateWeightsFromIPCReqOutput,
-                     self.update_weights_from_ipc_communicator.handle_recv,
-                 ),
-+                (
-+                    PostProcessWeightsReqOutput,
-+                    self.post_process_weights_communicator.handle_recv,
-+                ),
-                 (
-                     GetWeightsByNameReqOutput,
-                     self.get_weights_by_name_communicator.handle_recv,
-@@ -522,6 +531,17 @@ class TokenizerCommunicatorMixin:
- 
-         return success, message
- 
-+    async def post_process_weights(
-+        self: TokenizerManager,
-+        obj: PostProcessWeightsReqInput,
-+        request: Optional[fastapi.Request] = None,
-+    ) -> Tuple[bool, str]:
-+        """Trigger post-processing hooks for weights after loading (e.g., Marlin conversion)."""
-+        self.auto_create_handle_loop()
-+        async with self.model_update_lock.writer_lock:
-+            results = await self.post_process_weights_communicator(obj)
-+            return _Communicator.merge_results(results)
-+
-     async def init_weights_send_group_for_remote_instance(
-         self,
-         obj: InitWeightsSendGroupForRemoteInstanceReqInput,
-diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py
-index 0914a5230b..9114e3e713 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
-         async with self.is_pause_cond:
-             self.is_pause = True
-             if obj.mode != "abort":
--                await self.send_to_scheduler.send_pyobj(obj)
-+                self.send_to_scheduler.send_pyobj(obj)
-             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
-     async def continue_generation(self, obj: ContinueGenerationReqInput):
-         async with self.is_pause_cond:
-             self.is_pause = False
--            await self.send_to_scheduler.send_pyobj(obj)
-+            self.send_to_scheduler.send_pyobj(obj)
-             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
-         )
--        if (
--            state.first_token_time == 0.0
--            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()
-             state.last_completion_tokens = completion_tokens
--            self.metrics_collector.observe_time_to_first_token(
--                labels, state.first_token_time - state.created_time
--            )
-+            if self.disaggregation_mode != DisaggregationMode.PREFILL:
-+                self.metrics_collector.observe_time_to_first_token(
-+                    labels, state.first_token_time - state.created_time
-+                )
-         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
-                     num_new_tokens,
-                 )
-                 state.last_time = new_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
---- 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 (
-     InitWeightsUpdateGroupReqInput,
-     LoadLoRAAdapterFromTensorsReqInput,
-     LoadLoRAAdapterReqInput,
-+    PostProcessWeightsReqInput,
-     SendWeightsToRemoteInstanceReqInput,
-     UnloadLoRAAdapterReqInput,
-     UpdateWeightFromDiskReqInput,
-@@ -168,6 +169,11 @@ class BaseTpWorker(ABC):
-         success, message = self.model_runner.update_weights_from_ipc(recv_req)
-         return success, message
- 
-+    def post_process_weights(self, recv_req: PostProcessWeightsReqInput):
-+        """Perform optional post-processing on the updated model weights (e.g., Marlin conversion)."""
-+        success, message = self.model_runner.post_process_weights(recv_req)
-+        return success, message
-+
-     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
---- 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):
-             self._update_leaf_status(node)
-             self._update_host_leaf_status(node)
-             if node.parent is None:
--                assert (
--                    node is self.root_node
--                ), f"This request holds the node from another tree"
-+                # Node belongs to a stale (flushed) tree — stop traversal gracefully.
-+                break
-             node = node.parent
-         return delta
- 
-@@ -827,6 +826,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)
-+        self._update_host_leaf_status(node.parent)
-         return num_evicted
- 
-     def _evict_regular(self, node: TreeNode):
-@@ -1330,6 +1330,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)
-+                    self._update_host_leaf_status(node.parent)
-                 else:
-                     self._inc_hit_count(node, chunked)
-                     total_prefix_length += prefix_len
-@@ -1345,6 +1346,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)
-+                    self._update_host_leaf_status(new_node.parent)
-                 else:
-                     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
---- 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):
-         else:
-             assert self.page_size == 64
-         with (
--            torch.cuda.use_mem_pool(self.custom_mem_pool)
--            if self.custom_mem_pool
--            else nullcontext()
-+            (
-+                torch.cuda.use_mem_pool(self.custom_mem_pool)
-+                if self.custom_mem_pool
-+                else nullcontext()
-+            ),
-+            self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE),
-         ):
-             self.index_k_with_scale_buffer = [
-                 torch.zeros(
-@@ -1801,6 +1804,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._finalize_allocation_log(size)
- 
-     def get_index_k_with_scale_buffer(self, layer_id: int) -> torch.Tensor:
-@@ -1876,6 +1884,50 @@ class NSATokenToKVPool(MLATokenToKVPool):
-         ]
-         return data_ptrs, data_lens, item_lens
- 
-+    def get_cpu_copy(self, indices):
-+        # First, save the kv_buffer (inherited from MLATokenToKVPool)
-+        kv_cache_cpu = super().get_cpu_copy(indices)
-+
-+        # Additionally, save the index_k_with_scale_buffer (page-indexed)
-+        page_indices = indices[:: self.page_size] // self.page_size
-+        torch.cuda.synchronize()
-+        index_k_cpu = []
-+        chunk_size = self.cpu_offloading_chunk_size
-+        # Convert chunk_size from token-level to page-level
-+        page_chunk_size = max(1, chunk_size // self.page_size)
-+        for layer_id in range(self.layer_num):
-+            index_k_cpu.append([])
-+            for i in range(0, len(page_indices), page_chunk_size):
-+                chunk_page_indices = page_indices[i : i + page_chunk_size]
-+                idx_cpu = self.index_k_with_scale_buffer[layer_id][
-+                    chunk_page_indices
-+                ].to("cpu", non_blocking=True)
-+                index_k_cpu[-1].append(idx_cpu)
-+        torch.cuda.synchronize()
-+
-+        return {"kv": kv_cache_cpu, "index_k": index_k_cpu}
-+
-+    def load_cpu_copy(self, kv_cache_cpu_dict, indices):
-+        # Restore the kv_buffer (inherited from MLATokenToKVPool)
-+        super().load_cpu_copy(kv_cache_cpu_dict["kv"], indices)
-+
-+        # Restore the index_k_with_scale_buffer (page-indexed)
-+        page_indices = indices[:: self.page_size] // self.page_size
-+        index_k_cpu = kv_cache_cpu_dict["index_k"]
-+        torch.cuda.synchronize()
-+        chunk_size = self.cpu_offloading_chunk_size
-+        page_chunk_size = max(1, chunk_size // self.page_size)
-+        for layer_id in range(self.layer_num):
-+            for i in range(0, len(page_indices), page_chunk_size):
-+                chunk_page_indices = page_indices[i : i + page_chunk_size]
-+                idx_cpu = index_k_cpu[layer_id][i // page_chunk_size]
-+                assert idx_cpu.shape[0] == len(chunk_page_indices)
-+                idx_chunk = idx_cpu.to(
-+                    self.index_k_with_scale_buffer[0].device, non_blocking=True
-+                )
-+                self.index_k_with_scale_buffer[layer_id][chunk_page_indices] = idx_chunk
-+        torch.cuda.synchronize()
-+
-     def get_kv_size_bytes(self):
-         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
---- 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):
-         if self.disable:
-             return
- 
--        token_ids = req.fill_ids
-+        # Limit to kv_committed_len to avoid including tokens (e.g., the just-generated
-+        # token in disagg prefill) that don't have computed KV yet. If fill_ids is longer
-+        # than kv_committed_len, the extra tokens would produce stale values (0 from
-+        # req_to_token_pool initialization), leading to spurious tree nodes and memory
-+        # leak when page-aligned token counts happen to cross a page boundary.
-+        kv_committed_len = req.kv_committed_len
-+        token_ids = (
-+            req.fill_ids[:kv_committed_len]
-+            if kv_committed_len < len(req.fill_ids)
-+            else req.fill_ids
-+        )
-         kv_indices = self.req_to_token_pool.req_to_token[
-             req.req_pool_idx, : len(token_ids)
-         ]
-@@ -619,9 +629,8 @@ class RadixCache(BasePrefixCache):
-             node.lock_ref -= 1
-             self._update_leaf_status(node)
-             if node.parent is None:
--                assert (
--                    node is self.root_node
--                ), f"This request holds the node from another tree"
-+                # Node belongs to a stale (flushed) tree — stop traversal gracefully.
-+                break
-             node = node.parent
-         return 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
- 
--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
- 
-+    # 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
- 
-+    # --- 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
-+
-+        target_device = torch.device("cuda", torch.cuda.current_device())
-+
-+        if recv_req.restore_weights_before_load:
-+            for _, module in self.model.named_modules():
-+                quant_method = getattr(module, "quant_method", None)
-+
-+                # Check if the module supports restoring weights
-+                if quant_method is not None and hasattr(
-+                    quant_method, "restore_weights_before_loading"
-+                ):
-+
-+                    with device_loading_context(module, target_device):
-+                        quant_method.restore_weights_before_loading(module)
-+
-+        if recv_req.post_process_quantization:
-+            # Iterate through all modules to apply specific post-loading processing
-+            for _, module in self.model.named_modules():
-+                quant_method = getattr(module, "quant_method", None)
-+
-+                # Check if the module supports quantization post-processing
-+                if quant_method is not None and hasattr(
-+                    quant_method, "process_weights_after_loading"
-+                ):
-+
-+                    # Apply the post-processing (e.g., repacking weights for Marlin kernel)
-+                    with device_loading_context(module, target_device):
-+                        quant_method.process_weights_after_loading(module)
-+
-+        return True, "Success"
-+
- 
- 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
---- a/python/sglang/srt/models/glm4v_moe.py
-+++ b/python/sglang/srt/models/glm4v_moe.py
-@@ -52,11 +52,31 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-         self.num_fused_shared_experts = 0
-         self.determine_num_fused_shared_experts()
- 
--        self.model = Glm4MoeModel(
--            config,
--            quant_config,
--            prefix=add_prefix("language_model", prefix),
--        )
-+        if not self.config.encoder_only:
-+            self.model = Glm4MoeModel(
-+                config,
-+                quant_config,
-+                prefix=add_prefix("language_model", prefix),
-+            )
-+
-+            if self.pp_group.is_last_rank:
-+                if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
-+                    self.lm_head = self.model.embed_tokens
-+                else:
-+                    self.lm_head = ParallelLMHead(
-+                        config.vocab_size,
-+                        config.hidden_size,
-+                        quant_config=quant_config,
-+                        prefix=add_prefix("lm_head", prefix),
-+                        use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
-+                    )
-+            else:
-+                # ranks other than the last rank will have a placeholder layer
-+                self.lm_head = PPMissingLayer()
-+        else:
-+            # encoder_only mode: no language model, so no lm_head needed
-+            self.lm_head = None
-+
-         self.visual = Glm4vVisionModel(
-             config.vision_config,
-             quant_config=quant_config,
-@@ -64,24 +84,14 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-             use_data_parallel=self.use_data_parallel,
-         )
- 
--        if self.pp_group.is_last_rank:
--            if self.pp_group.world_size == 1 and self.config.tie_word_embeddings:
--                self.lm_head = self.model.embed_tokens
--            else:
--                self.lm_head = ParallelLMHead(
--                    config.vocab_size,
--                    config.hidden_size,
--                    quant_config=quant_config,
--                    prefix=add_prefix("lm_head", prefix),
--                    use_attn_tp_group=get_global_server_args().enable_dp_lm_head,
--                )
--        else:
--            # ranks other than the last rank will have a placeholder layer
--            self.lm_head = PPMissingLayer()
--
-         self.logits_processor = LogitsProcessor(config)
-         self.pooler = Pooler(pooling_type=PoolingType.LAST, normalize=True)
--        self.is_mrope_enabled = "mrope_section" in self.config.rope_scaling
-+        _rope_cfg = (
-+            getattr(self.config, "rope_scaling", None)
-+            or getattr(self.config, "rope_parameters", None)
-+            or {}
-+        )
-+        self.is_mrope_enabled = "mrope_section" in _rope_cfg
- 
-         # For EAGLE3 support
-         self.capture_aux_hidden_states = False
-@@ -219,6 +229,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                 # Skip loading extra bias for GPTQ models.
-                 if name.endswith(".bias") and name not in params_dict:
-                     continue
-+                # Skip loading visual/language model weights
-+                if (
-+                    self.config.encoder_only or self.config.language_only
-+                ) and name not in params_dict:
-+                    continue
-                 if name not in params_dict:
-                     continue
- 
-@@ -234,6 +249,8 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                     param_name, weight_name, expert_id, shard_id = mapping
-                     if weight_name not in name:
-                         continue
-+                    if "visual" in name or self.config.encoder_only:
-+                        continue
- 
-                     # Mark as expert weight regardless of whether we can process it
-                     is_expert_weight = True
-@@ -265,6 +282,11 @@ class Glm4vMoeForConditionalGeneration(Glm4vForConditionalGeneration):
-                     # Skip loading extra bias for GPTQ models.
-                     if name.endswith(".bias") and name not in params_dict:
-                         continue
-+                    # Skip loading mm/language parameters
-+                    if (
-+                        self.config.encoder_only or self.config.language_only
-+                    ) and name not in params_dict:
-+                        continue
-                     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
-+            )
- 
-         return hidden_states, residual
- 
-diff --git a/python/sglang/srt/models/qwen3_vl.py b/python/sglang/srt/models/qwen3_vl.py
-index d641826e33..3abc39ef32 100644
---- a/python/sglang/srt/models/qwen3_vl.py
-+++ b/python/sglang/srt/models/qwen3_vl.py
-@@ -711,14 +711,19 @@ class Qwen3LLMModel(Qwen3Model):
-                     hidden_states + residual if residual is not None else hidden_states
-                 )
- 
-+            deepstack_embeds = None
-+            if input_deepstack_embeds is not None:
-+                prev_layer_idx = layer_idx - 1
-+                if prev_layer_idx in self.deepstack_embed_to_decoder_layer:
-+                    sep = self.hidden_size * prev_layer_idx
-+                    deepstack_embeds = input_deepstack_embeds[
-+                        :, sep : sep + self.hidden_size
-+                    ]
-+
-             # SGLang applies residual at the START of the next layer, not at the END like HuggingFace.
-             # See: https://github.com/huggingface/transformers/blob/v5.0.0rc0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L549
-             # To match HF behavior, deepstack must be added AFTER residual: (hidden_states + residual) + deepstack
-             # The order matters because addition with different tensors is not associative in practice.
--            # Deepstack for prev_layer is applied at the start of current layer via post_residual_addition.
--            deepstack_embeds = self.get_deepstack_embeds(
--                layer_idx - 1, input_deepstack_embeds
--            )
-             hidden_states, residual = layer(
-                 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
---- a/python/sglang/srt/multimodal/processors/glm4v.py
-+++ b/python/sglang/srt/multimodal/processors/glm4v.py
-@@ -1,6 +1,9 @@
- 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.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):
-         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
-+        self.IM_START_TOKEN_ID = self.IMAGE_START_TOKEN_ID
-+        self.IM_END_TOKEN_ID = self.IMAGE_END_TOKEN_ID
- 
-         # Vision config
-         self.IMAGE_FACTOR = 28
-@@ -59,6 +64,36 @@ 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)
-+        mm_items = [
-+            MultimodalDataItem(
-+                modality=Modality.IMAGE,
-+                offsets=offsets,
-+                precomputed_embeddings=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),
-+            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]],
-diff --git a/python/sglang/srt/multimodal/processors/qwen_vl.py b/python/sglang/srt/multimodal/processors/qwen_vl.py
-index 4395654e4e..f9b5ea4abb 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):
-         **kwargs,
-     ):
-         entry_time = time.perf_counter()
--        base_output = self.load_mm_data(
-+        base_output = self.legacy_load_mm_data(
-             prompt=input_text,
-             image_data=image_data,
-             video_data=request_obj.video_data,
-diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
-index b080aeb168..5b29ebf566 100644
---- a/python/sglang/srt/server_args.py
-+++ b/python/sglang/srt/server_args.py
-@@ -635,6 +635,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:
-             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.",
-         )
-+        parser.add_argument(
-+            "--disable-indexer-rope-neox-style",
-+            action="store_true",
-+            help="Disable NSA indexer RoPE neox style (equivalent to INDEXER_ROPE_NEOX_STYLE=0). "
-+            "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",
-             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
---- 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_()
-         num_tokens = bs * self.num_tokens_per_bs
- 
-         # Common inputs
-@@ -350,8 +353,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_(
-+            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)
- 
-diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py
-index ac629c7ee5..c039d23508 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):
-             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)]
-+            if self.accept_length is not None:
-+                self.accept_length = self.accept_length[: len(new_indices)]
-+            if self.accept_length_cpu is not None:
-+                self.accept_length_cpu = self.accept_length_cpu[: len(new_indices)]
-         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):
-         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])
-+        if self.accept_length is not None and spec_info.accept_length is not None:
-+            self.accept_length = torch.cat(
-+                [self.accept_length, spec_info.accept_length]
-+            )
-+            self.accept_length_cpu = self.accept_length.tolist()
-+        elif self.accept_length is not None:
-+            zeros = torch.zeros(
-+                [spec_info.verified_id.shape[0]],
-+                dtype=self.accept_length.dtype,
-+                device=self.accept_length.device,
-+            )
-+            self.accept_length = torch.cat([self.accept_length, zeros])
-+            self.accept_length_cpu = self.accept_length.tolist()
-+        elif spec_info.accept_length is not None:
-+            zeros = torch.zeros(
-+                [self.verified_id.shape[0]],
-+                dtype=self.accept_length.dtype,
-+                device=self.accept_length.device,
-+            )
-+            self.accept_length = torch.cat([zeros, spec_info.accept_length])
-+            self.accept_length_cpu = self.accept_length.tolist()
- 
- 
- @dataclass
-diff --git a/python/sglang/srt/utils/common.py b/python/sglang/srt/utils/common.py
-index 4636128fa7..a9b61df393 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.",
-         "sglang.srt.layers.",
-         "sglang.srt.utils.",
-+        # --- slime ---
-+        "slime.",
-     }
- 
-     DENY_CLASSES = {
-diff --git a/python/sglang/srt/utils/weight_checker.py b/python/sglang/srt/utils/weight_checker.py
-index 3be16446e0..1b2371c839 100644
---- a/python/sglang/srt/utils/weight_checker.py
-+++ b/python/sglang/srt/utils/weight_checker.py
-@@ -69,6 +69,9 @@ def _check_tensors(
-         actual_should_compare,
-         actual,
-     ) in zip(expect_tensors, actual_tensors, strict=True):
-+        if ".cos_sin_cache" in expect_name:
-+            # skip cos/sin cache which is deterministic from shape and dtype and may have different shapes due to different implementations.
-+            continue
-         assert expect_name == actual_name, f"{expect_name=} {actual_name=}"
-         assert (
-             expect_should_compare == actual_should_compare
diff --git a/docker/version.txt b/docker/version.txt
index 54428a6001..dee5f8538b 100644
--- a/docker/version.txt
+++ b/docker/version.txt
@@ -1 +1 @@
-nightly-dev-20260430b
+nightly-dev-20260303a
diff --git a/docs/_static/image/arch.png b/docs/_static/image/arch.png
deleted file mode 100644
index b809bd3ffe..0000000000
Binary files a/docs/_static/image/arch.png and /dev/null differ
diff --git a/docs/_static/image/sglang_config.png b/docs/_static/image/sglang_config.png
deleted file mode 100644
index c4cde43ab2..0000000000
Binary files a/docs/_static/image/sglang_config.png and /dev/null differ
diff --git a/docs/_static/image/trace.png b/docs/_static/image/trace.png
deleted file mode 100644
index 847f4009cb..0000000000
Binary files a/docs/_static/image/trace.png and /dev/null differ
diff --git a/docs/en/advanced/megatron-config.md b/docs/en/advanced/megatron-config.md
deleted file mode 100644
index 973363ca81..0000000000
--- a/docs/en/advanced/megatron-config.md
+++ /dev/null
@@ -1,131 +0,0 @@
-# Megatron Config: Role-Based Training Overrides
-
-`--megatron-config-path` is a YAML-based configuration system for applying role-specific overrides on top of the shared Megatron CLI arguments. Today it is mainly intended for PPO actor / critic configuration.
-
-Unlike `--sglang-config`, `--megatron-config-path` does not manage deployment, routing, or GPU orchestration. Its only job is to decide which training arguments each role should finally use.
-
----
-
-## Design Overview
-
-By default, when `--megatron-config-path` is not used, both actor and critic inherit the Megatron / slime CLI arguments directly.
-
-With `--megatron-config-path`, the configuration is split into two layers:
-
-- **Shared CLI arguments** define the common Megatron topology, resource allocation, and default training parameters.
-- **Role-level YAML overrides** only specify the fields that should differ between actor and critic.
-
-**Key design principles:**
-
-- **CLI remains the shared baseline.** slime first parses the normal CLI arguments, then applies the YAML role overrides.
-- **Missing roles inherit automatically.** If a role is absent from the YAML file, it simply keeps the CLI arguments unchanged.
-- **Resource allocation is still controlled by CLI.** `num_nodes` and `num_gpus_per_node` in YAML are ignored; placement is still controlled by `--actor-num-*` / `--critic-num-*`.
-
----
-
-## Config Format
-
-The config file is a YAML document whose top-level `megatron` key contains a list of role entries:
-
-```yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-      save: /path/to/actor_ckpt
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-      save: /path/to/critic_ckpt
-```
-
-### Field Reference
-
-| Field | Type | Default | Description |
-|-------|------|---------|-------------|
-| `name` | `str` | Optional | Label for this entry. The runtime does not depend on it today, but keeping `default` is recommended for forward compatibility. |
-| `role` | `str` | **Required** | Role name. Currently supported values are `actor` and `critic`. |
-| `overrides` | `dict` | `{}` | Role-specific argument overrides applied on top of the shared CLI arguments. |
-| `args` | `dict` | `{}` | Backward-compatible alias for `overrides`. New configs should prefer `overrides`. |
-
-> **Note:** Keys inside `overrides` use argparse attribute names, not CLI flag names. For example, use `tensor_model_parallel_size` rather than `tensor-model-parallel-size`.
-
----
-
-## Usage Pattern
-
-A typical PPO setup looks like this:
-
-```yaml
-# megatron_ppo.yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-```bash
-python train.py \
-  --advantage-estimator ppo \
-  --use-critic \
-  --megatron-config-path megatron_ppo.yaml \
-  --tensor-model-parallel-size 2 \
-  --sequence-parallel \
-  --pipeline-model-parallel-size 1 \
-  --context-parallel-size 1 \
-  --expert-model-parallel-size 1 \
-  --expert-tensor-parallel-size 1 \
-  --actor-num-nodes 1 \
-  --actor-num-gpus-per-node 8 \
-  --critic-num-nodes 1 \
-  --critic-num-gpus-per-node 8 \
-  ...
-```
-
-In this setup:
-
-- CLI defines the shared topology and resource layout.
-- YAML defines the role-specific differences, such as `lr`, `load`, `save`, or optimizer / scheduler parameters.
-
-### Overriding Only One Role
-
-You can also override only one role and let the other inherit the shared CLI configuration. For example, changing only the critic learning rate:
-
-```yaml
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-In this case the actor keeps the shared CLI arguments unchanged.
-
----
-
-## Current Limitations
-
-- **PPO only for now.** `--megatron-config-path` is currently intended for PPO actor / critic role configuration. It is not the recommended interface for GRPO, REINFORCE++, and other critic-free workflows.
-- **Actor and critic must use the same Megatron parallel topology in current PPO.** In particular, topology-related settings such as `tensor_model_parallel_size`, `pipeline_model_parallel_size`, `context_parallel_size`, `expert_model_parallel_size`, `expert_tensor_parallel_size`, and `sequence_parallel` should not differ between actor and critic.
-- **Keep topology-related settings on CLI.** The safest current pattern is to keep parallelism and resource arguments in the shared CLI configuration, and only put role-specific differences in YAML, such as `lr`, `load`, `save`, warmup, and optimizer / scheduler settings.
-
-If you configure different parallel topologies for actor and critic, the behavior is currently unsupported and may fail during initialization or training.
-
----
-
-## FAQ
-
-### Q: Can I provide only an actor entry or only a critic entry?
-
-Yes. Missing roles automatically inherit the shared CLI arguments, so you do not need to duplicate everything in YAML.
-
-### Q: Can I move `--actor-num-nodes` or `--critic-num-gpus-per-node` into YAML?
-
-No. Resource allocation and placement groups are still controlled by CLI arguments, and the corresponding YAML fields are ignored.
\ No newline at end of file
diff --git a/docs/en/advanced/rfc-vllm-rollout-backend.md b/docs/en/advanced/rfc-vllm-rollout-backend.md
deleted file mode 100644
index 5c3cbd20f5..0000000000
--- a/docs/en/advanced/rfc-vllm-rollout-backend.md
+++ /dev/null
@@ -1,360 +0,0 @@
-# RFC: Add vLLM as a Rollout Backend in Slime
-
-- **Author**: \<your name\>
-- **Status**: Draft
-- **Audience**: Slime rollout/runtime maintainers, RL training maintainers
-- **Last Updated**: 2026-03-02
-
-## 1. Summary
-
-This RFC proposes adding **vLLM** as a first-class rollout backend in Slime while preserving current SGLang behavior and avoiding regressions in GRPO workflows.
-
-The design is based on:
-
-1. A backend-agnostic rollout request/response contract.
-2. Backend adapters (`SGLangClient`, `VLLMClient`) that isolate protocol differences.
-3. Capability-aware behavior for non-parity features (abort, routed experts, prompt logprobs).
-4. Managed vLLM mode from day one (Slime manages vLLM process lifecycle inside Ray, same as SGLang path).
-5. Weight sync via NCCL broadcast (GPU direct transfer, no disk I/O).
-
-## 2. Why this is needed
-
-Slime currently assumes SGLang behavior in multiple places:
-
-- rollout generation response schema (`meta_info.finish_reason.type`, `output_token_logprobs`, optional `routed_experts`)
-- router control-plane interactions (`/workers`, `/list_workers`, `/abort_request`)
-- rollout server startup flow in Ray
-
-Supporting vLLM is therefore not a URL replacement. It requires a compatibility layer across both data plane and control plane semantics.
-
-## 3. Goals and Non-goals
-
-### Goals
-
-- Add `--rollout-backend {sglang,vllm}`.
-- Keep SGLang path unchanged by default.
-- Keep trainer/algorithm interfaces stable.
-- Support GRPO rollout with explicit compatibility semantics and observability.
-
-### Non-goals (initial phase)
-
-- Full parity for all SGLang-specific features.
-- Colocate mode (training and rollout sharing same GPUs).
-- R3 routed expert replay on vLLM.
-- Multi-instance vLLM with router load balancing.
-
-## 4. Architecture and Interface Changes (Explicit)
-
-This section is the key communication point for Slime maintainers.
-
-### 4.1 Architecture changes
-
-#### End-to-end architecture (string diagram)
-
-```text
-                          +--------------------+
-                          |     TrainerLoop    |
-                          +---------+----------+
-                                    |
-                                    v
-                          +--------------------+
-                          |   RolloutFunction  |
-                          +---------+----------+
-                                    |
-                                    v
-                    +-------------------------------+
-                    | CanonicalRolloutRequest       |
-                    | (input_ids, sampling_params,  |
-                    |  return_logprob, prompt_text) |
-                    +---------------+---------------+
-                                    |
-                                    v
-                       +--------------------------+
-                       |   RolloutBackendClient   |
-                       +------------+-------------+
-                                    |
-                  +-----------------+-----------------+
-                  |                                   |
-                  v                                   v
-      +------------------------+          +------------------------+
-      |      SGLangClient      |          |       VLLMClient       |
-      +-----------+------------+          +------------+-----------+
-                  |                                    |
-                  v                                    v
-   +-------------------------------+      +------------------------------+
-   | SGLangRouter or SlimeRouter   |      | SlimeRouter(generic) or      |
-   | (SGLang control-plane aware)  |      | direct vLLM endpoint          |
-   +---------------+---------------+      +--------------+---------------+
-                   |                                     |
-                   v                                     v
-        +----------------------+             +----------------------+
-        |    SGLang workers    |             |    vLLM workers      |
-        +----------+-----------+             +----------+-----------+
-                   \                                     /
-                    \                                   /
-                     v                                 v
-                    +-----------------------------------+
-                    | CanonicalRolloutResponse          |
-                    | (text, token_ids, token_logprobs, |
-                    |  finish_reason, backend_raw)      |
-                    +----------------+------------------+
-                                     |
-                                     v
-                    +-----------------------------------+
-                    | SampleUpdate + Training Pipeline  |
-                    | (backend-agnostic consumption)    |
-                    +-----------------------------------+
-```
-
-#### Component responsibility map
-
-| Component | Responsibility before RFC | Responsibility after RFC |
-|---|---|---|
-| `RolloutFunction` | Contains generic rollout + SGLang protocol details | Contains generic rollout orchestration only |
-| Backend protocol layer | Implicit in rollout logic | Explicit via `RolloutBackendClient` adapters |
-| `SGLangClient` | N/A (scattered logic) | Owns SGLang request/response/control-plane specifics |
-| `VLLMClient` | N/A | Owns vLLM request/response mapping and retries |
-| Trainer/sample pipeline | Consumes SGLang-shaped fields indirectly | Consumes canonical fields, backend-agnostic |
-| Router integration | Mixed generic + SGLang-specific assumptions | SGLang-specific control-plane isolated to SGLang adapter |
-
-#### Control-plane behavior split
-
-| Control-plane behavior | SGLang path | vLLM initial path |
-|---|---|---|
-| Worker registration/discovery APIs | Supported | Not required in generic path |
-| Worker-level abort | Supported (`/abort_request`) | Fallback to timeout/cancel semantics |
-| Routed experts replay metadata | Supported | Explicitly unsupported (capability-gated) |
-| Health/load-balance routing | Existing SGLang/SlimeRouter behavior | SlimeRouter generic mode or direct endpoint |
-
-#### A) Rollout backend abstraction layer (new)
-
-- Introduce a backend client interface:
-  - `RolloutBackendClient`
-  - backend capability descriptor
-- Add concrete adapters:
-  - `SGLangClient` (existing behavior extraction)
-  - `VLLMClient` (new)
-
-**Impact**: rollout logic calls a unified backend interface instead of directly embedding SGLang HTTP semantics.
-
-#### B) Rollout execution path refactor
-
-- `sglang_rollout.generate` no longer directly depends on SGLang HTTP payload/response shape.
-- It builds a canonical request and consumes a canonical response.
-
-**Impact**: trainer-side logic remains stable while backend-specific details move into adapters.
-
-#### C) Startup path split in RolloutManager
-
-- Existing path (SGLang): keep current managed startup behavior.
-- vLLM path: managed mode -- Slime creates `VLLMEngine` Ray actors that launch and manage local vLLM server processes, just like SGLang path uses `SGLangEngine`.
-
-**Impact**: vLLM gets the same lifecycle management as SGLang (process startup, health check, shutdown).
-
-#### D) Weight sync: VLLMEngine with same interface as SGLangEngine
-
-Training-side weight updater (`UpdateWeightFromDistributed`) calls engine methods via Ray remote. The core call chain with source locations:
-
-```text
-Training side (Megatron actor, UNCHANGED)
-  UpdateWeightFromDistributed.update_weights()
-    -> engine.pause_generation.remote()                # Ray remote call
-        -> VLLMEngine.pause_generation()                # Ray actor method
-            -> requests.post("http://localhost:8000/sleep?level=2")
-            -> requests.post("http://localhost:8000/wake_up?tags=weights")
-
-    -> engine.flush_cache.remote()
-        -> VLLMEngine.flush_cache()                     # no-op, sleep level 2 covers this
-
-    -> engine.init_weights_update_group.remote()
-        -> VLLMEngine.init_weights_update_group()
-            -> requests.post("http://localhost:8000/collective_rpc",
-                 json={"method": "init_weight_update_group",
-                       "master_address": ..., "master_port": ...,
-                       "rank_offset": ..., "world_size": ...})
-
-    -> dist.broadcast(param, src=0, group=nccl_group)  # training process NCCL broadcast
-    -> engine.update_weights_from_distributed.remote()  # tell vLLM "receive these params"
-        -> VLLMEngine.update_weights_from_distributed()
-            -> requests.post("http://localhost:8000/collective_rpc",
-                 json={"method": "update_weight",
-                       "name": ..., "dtype_name": ..., "shape": ...})
-            -> vLLM worker: NCCL broadcast recv + model.load_weights()
-
-    -> engine.continue_generation.remote()
-        -> VLLMEngine.continue_generation()
-            -> requests.post("http://localhost:8000/wake_up?tags=kv_cache")
-
-
-Rollout generation side
-  sglang_rollout.generate()
-    -> VLLMClient.generate()
-        -> httpx.post("http://localhost:8000/v1/completions")  # direct to local vLLM
-```
-
-`VLLMEngine` exposes **the same method signatures** as `SGLangEngine`. Endpoint mapping:
-
-| Method (same signature) | SGLangEngine internal | VLLMEngine internal |
-|---|---|---|
-| `pause_generation()` | `POST /pause_generation` | `POST /sleep?level=2` + `POST /wake_up?tags=weights` |
-| `flush_cache()` | `GET /flush_cache` | no-op (sleep level 2 covers this) |
-| `init_weights_update_group(...)` | `POST /init_weights_update_group` | `POST /collective_rpc {"method":"init_weight_update_group",...}` |
-| `update_weights_from_distributed(...)` | `POST /update_weights_from_distributed` | `POST /collective_rpc {"method":"update_weight",...}` |
-| `continue_generation()` | `POST /continue_generation` | `POST /wake_up?tags=kv_cache` |
-
-**Impact**: training-side code (`UpdateWeightFromDistributed`, `train.py`, actor code) requires **zero changes**. Weight sync uses NCCL broadcast (GPU direct transfer), same efficiency as SGLang path.
-
-Source links:
-- [train.py#L88](../../../train.py#L88), [actor_group.py#L126](../../../slime/ray/actor_group.py#L126), [actor.py#L532](../../../slime/backends/megatron_utils/actor.py#L532)
-- [update_weight_from_distributed.py#L82](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L82), [#L89](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L89), [#L90](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L90), [#L280](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L280), [#L321](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L321), [#L139](../../../slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py#L139)
-- [sglang_engine.py#L415](../../../slime/backends/sglang_utils/sglang_engine.py#L415), [#L296](../../../slime/backends/sglang_utils/sglang_engine.py#L296), [#L373](../../../slime/backends/sglang_utils/sglang_engine.py#L373), [#L398](../../../slime/backends/sglang_utils/sglang_engine.py#L398), [#L420](../../../slime/backends/sglang_utils/sglang_engine.py#L420)
-- [sglang_rollout.py#L47](../../../slime/rollout/sglang_rollout.py#L47), [#L72](../../../slime/rollout/sglang_rollout.py#L72), [#L165](../../../slime/rollout/sglang_rollout.py#L165), [#L311](../../../slime/rollout/sglang_rollout.py#L311)
-- [rollout.py#L477](../../../slime/ray/rollout.py#L477), [#L1028](../../../slime/ray/rollout.py#L1028), [#L1041](../../../slime/ray/rollout.py#L1041)
-
-#### E) Router decision: no router for vLLM initial phase
-
-- SGLang Model Gateway: only supports SGLang workers, not applicable.
-- SlimeRouter: only needed for R3 / radix-tree caching; Qwen2.5-0.5B is not MoE and uses token-in/token-out.
-- Single vLLM instance, `VLLMClient` connects directly to local vLLM server port.
-
-### 4.2 Interface changes
-
-#### A) New CLI interfaces
-
-- `--rollout-backend {sglang,vllm}`
-- `--vllm-base-url`
-- `--vllm-api-mode` (e.g. OpenAI-compatible completion mode)
-- `--vllm-model`
-- `--vllm-max-retries`
-
-#### B) New internal canonical interfaces
-
-Add canonical rollout contract types:
-
-- `RolloutBackendRequest`
-- `RolloutBackendResponse`
-
-These carry backend-neutral fields such as:
-
-- input token ids
-- sampling params
-- output token ids/logprobs
-- canonical finish reason (`stop|length|abort`)
-- backend raw response for debugging
-
-#### C) Capability-gated interface behavior
-
-Backends declare capabilities (abort support, routed experts support, prompt logprobs support). Unsupported features are explicitly gated, logged, and/or failed fast.
-
-## 5. File-level Change Map (for maintainers)
-
-### New files
-
-- `slime/rollout/backends/base_client.py` -- backend client interface + capability model
-- `slime/rollout/backends/sglang_client.py` -- extracted SGLang rollout client
-- `slime/rollout/backends/vllm_client.py` -- vLLM rollout client
-- `slime/rollout/backends/__init__.py` -- adapter exports
-- `slime/backends/vllm_utils/vllm_engine.py` -- `VLLMEngine` Ray actor (analogous to `SGLangEngine`)
-- `slime/backends/vllm_utils/worker_extension.py` -- vLLM WorkerExtension for NCCL weight sync
-
-### Modified files
-
-- `slime/rollout/base_types.py` -- add canonical backend request/response types
-- `slime/rollout/sglang_rollout.py` -- use backend adapters instead of hardcoded SGLang protocol
-- `slime/utils/arguments.py` -- add `--rollout-backend`, vLLM args; skip SGLang parse when vLLM
-- `slime/ray/rollout.py` -- startup flow split: create `VLLMEngine` when `rollout_backend=vllm`
-
-### Unchanged files (by design)
-
-- `train.py` -- training loop does not know about backend
-- `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` -- engine method signatures are identical
-- `slime/backends/megatron_utils/actor.py` -- calls engine methods polymorphically
-
-## 6. Compatibility Matrix (initial)
-
-| Capability | SGLang | vLLM (initial) | Handling |
-|---|---|---|---|
-| Token-level response logprobs | Yes | Partial/endpoint-dependent | Adapter normalization |
-| Finish reason mapping (`stop|length|abort`) | Native | Different enums likely | Canonical mapping |
-| Worker-level abort | Yes | Typically no direct equivalent | timeout/cancel fallback strategy |
-| Prompt logprobs (OPD-related) | Yes | Partial/unknown by endpoint | capability-gated |
-| Routed experts replay (R3) | Yes | No | explicit unsupported gate |
-| SGLang worker API dependency | Yes | No | isolate in SGLang adapter |
-
-## 7. Phased Plan
-
-### Phase 1: Managed vLLM GRPO training (small-scale validation)
-
-**Goal**: Qwen2.5-0.5B GRPO 8-GPU sync training on GSM8K, loss/reward convergence comparable to SGLang.
-
-Key deliverables:
-- `VLLMEngine` Ray actor (process lifecycle, sleep/wake_up, weight sync via NCCL)
-- `VLLMClient` rollout adapter (request/response mapping, logprob/finish_reason normalization)
-- `RolloutBackendRequest/Response` canonical contract
-- `SGLangClient` extraction (isolate SGLang protocol details)
-- `start_rollout_servers` branching for vLLM path
-- Argument parsing split (`--rollout-backend vllm`)
-
-Technical decisions:
-- No router (single vLLM instance, direct connection)
-- Weight sync: NCCL broadcast via `VLLMEngine` with same method signatures as `SGLangEngine`
-- Non-colocate mode (separate training and rollout GPUs)
-- Rollout function: reuse `sglang_rollout.py` (backend-agnostic orchestration)
-- Training-side code: zero changes
-
-Acceptance criteria:
-- `num_rollout=3` smoke passes without errors
-- Weight sync correctness: rollout output changes with each training update
-- reward mean delta < 5% vs SGLang path (same seed/config, 20 steps)
-- Existing SGLang tests remain green
-
-### Phase 2: Performance, scale, and advanced features
-
-- Colocate mode support (GPU IPC weight transfer)
-- Multi-instance vLLM with load balancing
-- Abort/cancel strategy refinement
-- Prompt logprob support (OPD scenarios)
-- Deterministic computation verification
-- Larger model validation (e.g. Qwen3-4B)
-
-## 8. Validation strategy
-
-### Unit
-
-- finish reason normalization
-- token/logprob alignment
-- capability-gated behavior
-
-### Integration
-
-- SGLang vs vLLM response schema comparisons
-- timeout/retry/non-crash behavior checks
-
-### End-to-end GRPO
-
-- smoke (short rollout)
-- stability (medium horizon)
-- stress (higher concurrency / latency pressure)
-
-## 9. Risks and mitigations
-
-- **Logprob semantic mismatch**: strict adapter checks + canonicalization + tests.
-- **Abort mismatch**: capability model + timeout/cancel fallback + explicit logs.
-- **Training drift**: mandatory A/B runs with fixed seeds and aligned configs.
-- **Middleware assumptions**: enforce capability checks and default-disable incompatible middleware paths.
-
-## 10. Open questions for Slime maintainers
-
-1. Should vLLM initial path be direct endpoint only, or standardize via SlimeRouter generic mode immediately?
-2. What minimum logprob fidelity is required to claim GRPO support?
-3. Should OPD prompt-logprob paths be blocked for vLLM until full parity?
-4. What backend quality gates are required before default recommendation?
-
-## 11. Decision requested
-
-Approve incremental implementation with:
-
-- contract-first adapter architecture,
-- external vLLM initial support,
-- explicit capability-gated behavior,
-- strict SGLang non-regression requirement.
diff --git a/docs/en/advanced/sglang-config.md b/docs/en/advanced/sglang-config.md
deleted file mode 100644
index 70af72d2ec..0000000000
--- a/docs/en/advanced/sglang-config.md
+++ /dev/null
@@ -1,453 +0,0 @@
-# SGLang Config: Advanced Engine Deployment
-
-`--sglang-config` is a YAML-based configuration system for fine-grained control over SGLang engine deployment in slime. It enables **multi-model serving**, **Prefill-Decode (PD) disaggregation**, **heterogeneous server groups**, and can even serve as a **standalone SGLang launcher** for complex inference topologies.
-
----
-
-## Architecture Overview
-
-In the default setup (without `--sglang-config`), slime deploys a single model behind a single router with uniform server groups:
-
-![architecture overview](../../_static/image/arch.png)
-
-With `--sglang-config`, the SGLang deployment expands into a multi-model, multi-router topology:
-
-![sglang-config architecture](../../_static/image/sglang_config.png)
-
-**Key design principles:**
-
-- **Each model gets its own router.** Models are isolated at the routing layer, allowing independent load balancing and fault tolerance.
-- **Server groups within a model can be heterogeneous.** Different groups can have different TP sizes, worker types (prefill/decode/regular), and SGLang server argument overrides.
-- **Weight sync is per-model.** Only models with `update_weights: true` receive weight updates from training. Frozen models (reference, reward, etc.) are served as-is.
-
----
-
-## Config Format
-
-The config file is a YAML document with a top-level `sglang` key containing a list of model definitions:
-
-```yaml
-sglang:
-  - name: <model_name>              # Required. Unique identifier for this model.
-    model_path: <path>              # Optional. HF checkpoint path. Defaults to --hf-checkpoint.
-    update_weights: <bool>          # Optional. Whether to sync weights from training. Auto-inferred.
-    num_gpus_per_engine: <int>      # Optional. Default TP size for all groups in this model.
-    server_groups:                  # Required. List of server group configurations.
-      - worker_type: <type>         # Required. One of: regular, prefill, decode, placeholder.
-        num_gpus: <int>             # Required. Total GPUs allocated to this group.
-        num_gpus_per_engine: <int>  # Optional. TP size override for this group.
-        overrides: <dict>           # Optional. SGLang ServerArgs field overrides.
-```
-
-### Field Reference
-
-#### Model-Level Fields
-
-| Field | Type | Default | Description |
-|-------|------|---------|-------------|
-| `name` | `str` | **Required** | Unique name for this model (e.g., `"actor"`, `"ref"`, `"reward"`). Used as the key in `args.sglang_model_routers`. |
-| `model_path` | `str` | `args.hf_checkpoint` | HuggingFace checkpoint path. All server groups within a model must use the same model path. |
-| `update_weights` | `bool` | Auto | Whether this model receives weight updates from training. When not set, automatically inferred: `true` if `model_path` matches `--hf-checkpoint`, `false` otherwise. |
-| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | Default TP size for server groups in this model. Individual groups can override. |
-| `server_groups` | `list` | **Required** | List of `ServerGroupConfig` entries defining the engine topology. (`engine_groups` is accepted as a backward-compatible alias.) |
-
-#### Server Group Fields
-
-| Field | Type | Default | Description |
-|-------|------|---------|-------------|
-| `worker_type` | `str` | **Required** | Engine type: `regular` (standard), `prefill` (PD prefill worker), `decode` (PD decode worker), or `placeholder` (reserve GPU slots without launching engines). |
-| `num_gpus` | `int` | **Required** | Total number of GPUs for this group. Must be > 0. |
-| `num_gpus_per_engine` | `int` | Model's `num_gpus_per_engine` | TP size override. Number of GPUs per engine instance. |
-| `overrides` | `dict` | `{}` | SGLang `ServerArgs` field overrides. Applied on top of `--sglang-*` CLI args with highest priority. |
-
-### Worker Types
-
-| Type | Description | Use Case |
-|------|-------------|----------|
-| `regular` | Standard SGLang engine | Default mode, handles both prefill and decode |
-| `prefill` | PD disaggregation prefill worker | Dedicated to prompt processing; paired with `decode` workers |
-| `decode` | PD disaggregation decode worker | Dedicated to token generation; paired with `prefill` workers |
-| `placeholder` | Reserves GPU slots, no engine created | Reserve GPUs for training co-location or future use |
-
----
-
-## Usage Patterns
-
-### 1. Basic Single-Model Deployment
-
-The simplest config replicates the default behavior:
-
-```yaml
-# sglang_basic.yaml
-sglang:
-  - name: default
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_basic.yaml \
-  --rollout-num-gpus 8 \
-  --rollout-num-gpus-per-engine 2 \
-  ...
-```
-
-This creates 4 engines (8 GPUs ÷ 2 GPUs/engine) behind a single router.
-
-### 2. PD Disaggregation
-
-Separate prefill and decode phases onto dedicated server groups for better throughput in multi-turn and agentic workloads:
-
-```yaml
-# sglang_pd.yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2    # 2 prefill engines, TP=2
-      - worker_type: decode
-        num_gpus: 12
-        num_gpus_per_engine: 4    # 3 decode engines, TP=4
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_pd.yaml \
-  --rollout-num-gpus 16 \
-  ...
-```
-
-**Why PD disaggregation?** In multi-turn scenarios, prefill and decode have different compute characteristics. Prefill is compute-bound (processes the entire prompt), while decode is memory-bandwidth-bound (generates tokens one-by-one). Disaggregating them allows:
-- Using smaller TP for prefill (higher throughput per GPU)
-- Using larger TP for decode (lower latency)
-- Independent scaling of prefill vs. decode capacity
-
-> **Note:** PD disaggregation uses SGLang Model Gateway (sgl-router) with `pd_disaggregation=True`.
-
-### 3. Multi-Model Serving
-
-Deploy multiple models simultaneously, each behind its own router:
-
-```yaml
-# sglang_multi_model.yaml
-sglang:
-  - name: actor
-    update_weights: true              # receives weight updates from training
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: ref
-    model_path: /path/to/ref_model    # different model checkpoint
-    update_weights: false              # frozen, no weight updates
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-
-  - name: reward
-    model_path: /path/to/reward_model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_multi_model.yaml \
-  --rollout-num-gpus 16 \
-  --hf-checkpoint /path/to/actor_model \
-  --rollout-function-path my_rollout.generate_rollout \
-  ...
-```
-
-**Accessing models in custom rollout functions:**
-
-```python
-from slime.rollout.sglang_rollout import get_model_url
-from slime.utils.http_utils import post
-
-async def my_generate(args, sample, sampling_params):
-    # Route to the actor model (default)
-    actor_url = get_model_url(args, "actor", "/generate")
-    output = await post(actor_url, {"text": sample.prompt, "sampling_params": sampling_params})
-    
-    # Route to the reference model
-    ref_url = get_model_url(args, "ref", "/generate")
-    ref_output = await post(ref_url, {"text": sample.prompt, "sampling_params": sampling_params})
-    
-    # Route to the reward model (e.g., OpenAI-compatible API)
-    reward_url = get_model_url(args, "reward", "/v1/chat/completions")
-    reward_output = await post(reward_url, {...})
-    
-    ...
-```
-
-The `get_model_url()` helper reads from `args.sglang_model_routers`, a dict mapping model names to `(ip, port)` tuples that is automatically populated after engine startup.
-
-### 4. Multi-Model with PD Disaggregation
-
-Combine multi-model and PD disaggregation for maximum flexibility:
-
-```yaml
-# sglang_full.yaml
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2
-      - worker_type: decode
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: ref
-    model_path: /path/to/ref_model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-### 5. Placeholder Groups for GPU Reservation
-
-Use `placeholder` groups to reserve GPU slots without creating engines. This is useful for co-located training where some GPUs need to be reserved for training:
-
-```yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: regular
-        num_gpus: 6
-        num_gpus_per_engine: 2
-      - worker_type: placeholder
-        num_gpus: 2                   # reserve 2 GPUs (no engines created)
-```
-
-### 6. Per-Group ServerArgs Overrides
-
-Use `overrides` to apply SGLang `ServerArgs` fields to specific server groups without affecting others:
-
-```yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-        overrides:
-          mem_fraction_static: 0.85
-          context_length: 32768
-          chunked_prefill_size: 4096
-          enable_torch_compile: true
-```
-
-Overrides take **highest priority**, overriding both the base `--sglang-*` CLI args and model-level defaults. This is especially useful for:
-- Different memory configurations per group
-- Different context lengths for prefill vs. decode
-- Enabling experimental features on specific groups
-
-### 7. Standalone SGLang Launcher
-
-While `--sglang-config` is designed for slime's training pipeline, it also works as a powerful launcher for pure inference scenarios using the `--rollout-external` pattern or by configuring slime to focus solely on serving.
-
-**Using external engines with a pre-launched topology:**
-
-For complex production deployments, you may want to pre-launch SGLang engines independently and connect them to slime:
-
-```bash
-# Step 1: Launch SGLang engines externally
-python -m sglang.launch_server --model-path /path/to/model --port 10090 ...
-python -m sglang.launch_server --model-path /path/to/model --port 10091 ...
-
-# Step 2: Connect slime to external engines
-python train.py \
-  --rollout-external \
-  --rollout-external-engine-addrs host1:10090 host2:10091 \
-  ...
-```
-
-> **Note:** `--sglang-config` and `--rollout-external` are mutually exclusive. Use `--sglang-config` when you want slime to manage the full engine lifecycle; use `--rollout-external` when engines are pre-deployed.
-
----
-
-## Router Configuration
-
-Each model in the config gets its own independent router (SGLang Model Gateway by default).
-
-### Router Policies
-
-You can configure the routing policy:
-
-```bash
---router-policy round_robin        # Simple round-robin
---router-policy consistent_hashing  # Session affinity for multi-turn
---router-policy cache_aware         # Cache-aware routing (default)
-```
-
-### Session-Affinity Routing for Multi-Turn Agents
-
-For multi-turn dialogues and agentic workloads, session affinity ensures that all requests belonging to the same conversation are routed to the same backend worker. This significantly improves prefix cache hit rates because the worker already has the conversation history cached.
-
-slime automatically assigns each sample a unique `session_id` (stored in `sample.session_id`). When the router policy is `consistent_hashing`, this ID is passed as the `X-SMG-Routing-Key` header, and SGLang Model Gateway uses it to deterministically route all turns of the same session to the same worker.
-
-```bash
---router-policy consistent_hashing
-```
-
-**How it works:**
-
-1. Each sample is assigned a unique `session_id` via UUID
-2. On each request, slime passes `X-SMG-Routing-Key: <session_id>` in the HTTP header
-3. SGLang Model Gateway's consistent hashing policy maps this key to a specific worker
-4. Subsequent turns reuse the same `session_id`, ensuring they hit the same worker
-
----
-
-## Resolution Rules
-
-When the config is loaded, slime applies the following resolution cascade:
-
-1. **GPU per engine fallback:** Group `num_gpus_per_engine` → Model `num_gpus_per_engine` → `args.rollout_num_gpus_per_engine`
-2. **Model path fallback:** Group `overrides.model_path` → Model `model_path` → `args.hf_checkpoint`
-3. **Weight update inference:** If `update_weights` is not set:
-   - `true` if the effective model path matches `--hf-checkpoint`
-   - `false` otherwise (with a warning)
-4. **Total GPU validation:** The sum of all `num_gpus` across all groups across all models must equal `--rollout-num-gpus`.
-
----
-
-## Mutual Exclusion
-
-`--sglang-config` is mutually exclusive with:
-
-| Flag | Conflict Reason |
-|------|----------------|
-| `--prefill-num-servers` | PD disaggregation is configured via `server_groups` in the YAML |
-| `--rollout-external` | External engines have their own topology; config manages the lifecycle internally |
-
----
-
-## Complete Example: Multi-Model Agentic Training
-
-Below is a complete example showing a multi-model setup for agentic RL training with PD disaggregation on 32 GPUs:
-
-**Config file (`sglang_agent.yaml`):**
-
-```yaml
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2
-        overrides:
-          chunked_prefill_size: 8192
-      - worker_type: decode
-        num_gpus: 12
-        num_gpus_per_engine: 4
-        overrides:
-          mem_fraction_static: 0.88
-
-  - name: ref
-    model_path: /data/models/Qwen3-32B
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: reward
-    model_path: /data/models/reward-model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-```
-
-**Launch command:**
-
-```bash
-python train.py \
-  --sglang-config sglang_agent.yaml \
-  --hf-checkpoint /data/models/Qwen3-8B \
-  --rollout-num-gpus 32 \
-  --rollout-function-path my_agent.rollout.generate_rollout \
-  --custom-rm-path my_agent.reward.reward_func \
-  --advantage-estimator grpo \
-  --n-samples-per-prompt 8 \
-  ...
-```
-
-**Custom rollout function (`my_agent/rollout.py`):**
-
-```python
-from slime.rollout.sglang_rollout import get_model_url
-from slime.utils.http_utils import post
-
-async def generate_with_models(args, sample, sampling_params):
-    """Generate using actor, score with reward model, compare with reference."""
-    
-    # Generate from actor
-    actor_url = get_model_url(args, "actor", "/generate")
-    actor_output = await post(actor_url, {
-        "text": sample.prompt,
-        "sampling_params": sampling_params,
-        "return_logprob": True,
-    })
-    
-    # Get reference logprobs for KL penalty
-    ref_url = get_model_url(args, "ref", "/generate")
-    ref_output = await post(ref_url, {
-        "text": sample.prompt + actor_output["text"],
-        "sampling_params": {"max_new_tokens": 0, "temperature": 0},
-        "return_logprob": True,
-    })
-    
-    # Score with reward model
-    reward_url = get_model_url(args, "reward", "/v1/chat/completions")
-    reward_output = await post(reward_url, {
-        "model": "reward",
-        "messages": [{"role": "user", "content": sample.prompt + actor_output["text"]}],
-    })
-    
-    # ... process outputs and return Sample
-```
-
----
-
-## FAQ
-
-### Q: Can I mix PD and regular groups in the same model?
-
-No. PD disaggregation requires that a model's server groups are either all prefill/decode pairs or all regular. Mixing `regular` with `prefill`/`decode` in the same model is not supported.
-
-### Q: What happens if `num_gpus` is not divisible by `num_gpus_per_engine`?
-
-For multi-node engines (where `num_gpus_per_engine > num_gpus_per_node`), the division is based on the local GPU count per node. For example, with 8 GPUs/node and `num_gpus_per_engine: 16`, each engine spans 2 nodes.
-
-### Q: Can different server groups within a model use different model paths?
-
-No. All server groups within a model must share the same `model_path`. This is validated during `resolve()`. If you need different models, define them as separate model entries.
-
-### Q: How do I access the router address for a specific model at runtime?
-
-Use `get_model_url(args, "model_name", "/endpoint")` from `slime.rollout.sglang_rollout`. It reads from `args.sglang_model_routers`, which is a dict `{ model_name: (ip, port) }` populated automatically.
-
-### Q: Can I use `--sglang-config` without training (inference only)?
-
-While `--sglang-config` is designed for slime's training loop, you can effectively use it for inference-only scenarios by configuring a rollout-only run. For fully standalone SGLang serving, consider using SGLang's native `launch_server` directly or the `--rollout-external` mode for connecting to pre-deployed engines.
-
-### Q: What is the relationship between `--sglang-config` and `--prefill-num-servers`?
-
-`--prefill-num-servers` is the legacy way to enable PD disaggregation (it creates a single model with prefill + decode groups). `--sglang-config` is the newer, more flexible approach. They are mutually exclusive. We recommend migrating to `--sglang-config` for all new deployments.
diff --git a/docs/en/advanced/slime-router.md b/docs/en/advanced/slime-router.md
new file mode 100644
index 0000000000..f4d3ba3099
--- /dev/null
+++ b/docs/en/advanced/slime-router.md
@@ -0,0 +1,138 @@
+# Slime Router
+
+Slime includes an optional SlimeRouter used during rollout / data generation. It is a lightweight HTTP router/proxy that sits in front of one or more SGLang worker servers and adds training-oriented capabilities that are not the main goal of serving-focused routers.
+
+---
+
+## 1. What is SlimeRouter?
+
+SlimeRouter is a small FastAPI service that:
+
+- Registers workers (SGLang HTTP servers) into a local pool
+- Routes requests to a selected worker (simple least-inflight load balancing)
+- Proxies arbitrary paths to the selected worker (e.g. `/generate`)
+- Runs periodic health checks and quarantines unhealthy workers
+- Supports middleware plugins (via `--slime-router-middleware-paths`) to implement rollout-specific processing (e.g. caching, request/response transforms)
+
+In slime's architecture, the router is part of the rollout system ("SGLang + router") that generates samples and pushes them into the data buffer.
+
+### How it is launched
+
+In distributed training, slime will start a router automatically when `--sglang-router-ip` is not provided:
+
+- If `--use-slime-router` is set, slime starts SlimeRouter
+- Otherwise, slime starts SGLang Model Gateway
+
+---
+
+## 2. Why we need SlimeRouter
+
+Unlike production inference, RL rollout needs to capture additional metadata for training: token-level logprobs, loss masks, and (for MoE models) expert routing decisions. SlimeRouter provides these capabilities through its middleware system and passthrough proxy design.
+
+### 2.1 Radix-tree cache (transparent token management)
+
+> Use this when your rollout pipeline is text-in/text-out and you cannot reliably persist token IDs; if you already control token-in/token-out (e.g. search r1, multiturn VLM examples), you likely don't need the radix-tree cache.
+
+Text-in text-out interfaces can cause token retokenization mismatches - re-tokenizing text at training time may produce different token sequences than rollout, breaking per-token alignment needed for PPO/GRPO losses.
+
+The radix-tree cache solves this transparently: it intercepts text-based requests, tokenizes them, and stores trajectories (text, token IDs, logprobs, loss masks) keyed by the text prefix. After rollout finishes, calling `/retrieve_from_text` returns the exact token sequence with aligned metadata, without requiring any changes to your rollout code.
+
+Implementation-wise, the radix-tree cache:
+
+- Accepts text plus tokens/metadata and stores them in a radix tree
+- Uses longest-prefix matching to reuse cached token sequences (enabling token-in/token-out downstream)
+- Allows insertion of new text continuations as rollout proceeds (multiple trajectories per prompt, e.g. GRPO)
+- Periodically cleans up stale nodes to control memory usage
+
+Use the radix cache when you have text-based rollout code and want token-level precision without rewriting, or when running GRPO with multiple trajectories sharing the same prompt prefix.
+
+### 2.2 Rollout routing replay (R3) for MoE
+
+For MoE models, slime supports rollout routing replay (R3): record expert routing decisions during rollout and replay them during training to improve stability.
+
+#### SGLang side
+
+SGLang provides expert routing capture via:
+
+- `--enable-return-routed-experts`: server argument to enable routing capture
+- `RoutedExpertsCapturer`: captures `topk_ids` (selected expert IDs) at each MoE layer during forward pass
+- `return_routed_experts`: request parameter to retrieve routing data
+- Returns `routed_experts` in response `meta_info` - a `[seq_len - 1, num_layers, top_k]` tensor of expert IDs
+
+#### Slime side
+
+Slime consumes the routing data and replays it during training:
+
+- `--use-slime-router --use-rollout-routing-replay`: both flags required to enable R3
+- Rollout sends `return_routed_experts=True` and stores results in `sample.rollout_routed_experts`
+- Training calls `fill_routing_replay()` to load routing data into `RoutingReplay` objects
+- During forward pass, recorded routing decisions are replayed instead of recomputed
+
+#### Why SlimeRouter is needed
+
+We need SlimeRouter because the SGLang worker returns routed experts in the response (`meta_info.routed_experts`) when the request sets `return_routed_experts=true`, and SlimeRouter preserves this field end-to-end. SGLang Model Gateway may drop this extra metadata when it reconstructs responses with a fixed schema (see section 3).
+
+---
+
+## 3. Differences vs SGLang Model Gateway
+
+SlimeRouter and SGLang Model Gateway can both route requests to workers, but they are optimized for different goals.
+
+### Key differences
+
+SlimeRouter is a lightweight Python/FastAPI proxy that acts as a passthrough to SGLang workers. This passthrough design enables RL-specific features like radix-tree trajectory caching and R3 (which require preserving raw response metadata like `routed_experts`).
+
+SGLang Model Gateway is a high-performance Rust-based router optimized for large-scale inference: async non-blocking routing, advanced fault tolerance (retries, circuit breakers), multiple load balancing policies (including cache-aware routing), and PD disaggregation support. However, it reconstructs responses with a fixed schema, so it does not preserve the metadata needed for slime's R3 flow.
+
+For more details on SGLang Model Gateway, see the [official documentation](https://docs.sglang.io/advanced_features/sgl_model_gateway.html).
+
+### When to use which
+
+- Use SlimeRouter when you need R3 or radix-tree caching
+- Use SGLang Model Gateway for everything else (recommended default)
+
+---
+
+## 4. Session-Affinity Routing for Multi-Turn Agents
+
+When using SGLang Model Gateway with consistent hashing routing policy, Slime automatically assigns each rollout session a unique session ID and uses it as the routing key to enable session affinity.
+
+### What is session affinity?
+
+Session affinity (also called sticky sessions) ensures that all requests belonging to the same conversation or agent session are routed to the same backend worker. This is beneficial for:
+
+- **Multi-turn dialogues**: Keeping the same worker improves prefix cache hit rates
+- **Multi-agent systems**: Ensures agent state consistency and better resource locality
+- **Debugging**: Makes it easier to trace and debug specific sessions
+
+### How it works
+
+When the rollout system generates samples, each sample is assigned a unique `session_id`. This ID is:
+
+1. Automatically generated using UUID for each sample
+2. Stored in `sample.session_id` field
+3. Passed as `X-SMG-Routing-Key` header when the router policy is `consistent_hashing`
+
+The SGLang Model Gateway's consistent hashing policy then uses this routing key to deterministically select the same worker for all requests with the same session ID.
+
+### Configuration
+
+To enable session-affinity routing, simply configure the router policy in Slime:
+
+```bash
+--sglang-router-policy consistent_hashing
+```
+
+Slime will automatically start SGLang Model Gateway with the consistent hashing policy.
+
+> **Note**: If you encounter an error about the `consistent_hashing` policy not being available, upgrade sglang-router:
+> ```bash
+> pip install -U sglang-router
+> ```
+
+### Notes
+
+- Each sample gets its own unique session ID
+- Different samples in the same group may be routed to different workers
+- The same sample's subsequent turns will maintain the same session ID
+- Currently, this feature is only available for SGLang Model Gateway
diff --git a/docs/en/blogs/introducing_slime.md b/docs/en/blogs/introducing_slime.md
index 06eb60f20c..f35c7d65b2 100644
--- a/docs/en/blogs/introducing_slime.md
+++ b/docs/en/blogs/introducing_slime.md
@@ -36,7 +36,7 @@ It wasn’t always like this: no one forks PyTorch just for a new dataloader. We
 
 slime views the data sampling in RL differently. We manage all SGLang servers within slime with [sgl-router](https://github.com/sgl-project/sglang/tree/main/sgl-router) and provide an interface for the data generation component, **allowing users to inject custom logic and freely interact with SGLang servers**. Unleash their creativity.
 
-![slime architecture](../../_static/image/arch.png)
+![slime architecture](/imgs/arch.png)
 
 With the sgl-router, users only need to send HTTP requests to a single endpoint. By exposing this endpoint, complex agent environments can directly interact with slime through an OpenAI-compatible API — no need to modify the environment, and training-deployment consistency is preserved.
 
diff --git a/docs/en/developer_guide/ci.md b/docs/en/developer_guide/ci.md
index f84f46879a..7e93f38291 100644
--- a/docs/en/developer_guide/ci.md
+++ b/docs/en/developer_guide/ci.md
@@ -6,7 +6,7 @@ slime uses GitHub Actions for CI. Tests are triggered by **PR labels** — addin
 
 The workflow is defined in `.github/workflows/pr-test.yml` (auto-generated from `pr-test.yml.j2`). Each CI job:
 
-1. Runs on a self-hosted GPU runner via `docker run`; most tests use `slimerl/slime:latest`, while image validation uses `slimerl/slime-test:latest`.
+1. Runs on a self-hosted GPU runner inside a Docker container (`slimerl/slime:latest`).
 2. Installs slime with `pip install -e . --no-deps`.
 3. Acquires the required GPUs via `tests/ci/gpu_lock_exec.py --count <num_gpus>`.
 4. Executes the test file: `python <test_path>.py` or `python tests/<test_file>.py`, depending on whether the test lives under `tests/` or a subdirectory such as `tests/plugin_contracts/`.
@@ -20,8 +20,9 @@ Add a label to your PR to trigger the corresponding test suite:
 | Label | Job | Description |
 |---|---|---|
 | `run-ci-short` | `e2e-test-short` | Lightweight smoke tests with Qwen2.5-0.5B (4 GPUs). Fast feedback loop. |
+| `run-ci-fsdp` | `e2e-test-fsdp` | FSDP backend tests (true on-policy, VL, megatron-fsdp alignment). |
 | `run-ci-megatron` | `e2e-test-megatron` | Core Megatron training tests covering dense, MoE, PPO, MTP, OPD, etc. |
-| `run-ci-precision` | `e2e-test-precision` | Numerical precision validation (parallel check). |
+| `run-ci-precision` | `e2e-test-precision` | Numerical precision validation (parallel check, megatron-fsdp alignment). |
 | `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint save/load correctness (sync and async-save). |
 | `run-ci-image` | `e2e-test-image` | Full test suite run on `slimerl/slime-test:latest` image (for image validation). |
 | `run-ci-changed` | `e2e-test-changed` | **Dynamically** detects new/modified test files in the PR and runs only those. |
@@ -56,7 +57,7 @@ Since this includes every test, it consumes significant GPU time — use it spar
 This is the primary label for validating Megatron-backend changes. It covers:
 
 - Dense models: GLM4-9B, Qwen3-4B (PPO)
-- MoE models: Qwen3-30B-A3B (with DeepEP + FP8), Qwen3.6-35B-A3B PD + Mooncake, Moonlight-16B-A3B
+- MoE models: Qwen3-30B-A3B (with/without DeepEP + FP8), Moonlight-16B-A3B
 - Specialized: MiMo-7B MTP, Qwen2.5-0.5B debug rollout-then-train, OPD with sglang teacher
 
 All tests use 8 GPUs. If you are modifying Megatron training logic, loss computation, or checkpoint conversion, this is the label to use.
@@ -75,7 +76,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"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     # Download datasets as needed ...
 
 def execute():
diff --git a/docs/en/developer_guide/trace.md b/docs/en/developer_guide/trace.md
deleted file mode 100644
index 6fb0c6967b..0000000000
--- a/docs/en/developer_guide/trace.md
+++ /dev/null
@@ -1,119 +0,0 @@
-# Trace Viewer
-
-slime can attach lightweight execution traces to each rollout sample. These traces capture span-style events such as generation and reward-model calls, and they can be inspected later from a saved rollout debug dump.
-
-![trace timeline viewer](../../_static/image/trace.png)
-
-## Save rollout trace data
-
-To inspect traces later, save rollout debug data during a run:
-
-```bash
-python train.py \
-    ... \
-    --save-debug-rollout-data /path/to/debug/rollout_{rollout_id}.pt
-```
-
-Each saved `.pt` file contains the rollout samples together with their `trace` payloads. You can also replay the same dump later with `--load-debug-rollout-data`.
-
-## Open the timeline viewer
-
-Use the trace viewer script on a saved rollout dump:
-
-```bash
-python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt
-```
-
-The script generates:
-
-- `rollout_0.trace_timeline_cache.json`
-- `rollout_0.trace_timeline_viewer.html`
-
-By default it also starts a local static server so you can open the generated HTML immediately. If you only want the files, use `--no-serve`.
-
-## How to read the viewer
-
-- Each row corresponds to one sample.
-- Bars represent spans, while point markers represent instant events.
-- Span attributes recorded at the start or end of `trace_span(...)` are shown in the details panel.
-- When SGLang returns PD disaggregation timings, the viewer adds synthetic `[P]` and `[D]` lanes to break out prefill/decode work.
-- When PD is not enabled, those virtual lanes are omitted automatically and the base trace still renders normally.
-
-## Instrument custom code
-
-For custom rollout or reward code, reuse helpers from `slime.utils.trace_utils`:
-
-- `trace_span(target, name, attrs=...)`: record a duration span.
-- `trace_event(target, name, attrs=...)`: record an instant event.
-- `trace_function(name, ...)`: wrap a whole sync/async function in a span.
-- `bind_trace(sample)`: ensure a sample already has a trace carrier before passing it across helpers or tasks.
-
-### `trace_span` vs `trace_function`
-
-Use `trace_span(...)` when you only want to trace part of a function body, or when you need to update end-of-span attrs from inside the block.
-
-Use `trace_function(...)` when the whole function should be represented as one span. Internally it resolves the trace target and then opens a `trace_span(...)` around the function call, so it works for both sync and async functions.
-
-The decorator is what slime uses for the main rollout pipeline. For example, `generate_and_rm(...)` is traced per sample and `generate_and_rm_group(...)` is traced per sample group:
-
-```python
-from slime.utils.trace_utils import trace_function
-
-
-@trace_function("generate_and_rm", target="sample")
-async def generate_and_rm(args, sample, sampling_params, evaluation=False):
-    ...
-
-
-@trace_function(
-    "generate_and_rm_group",
-    target="group",
-    attrs_getter=lambda args, group, sampling_params, evaluation=False: {"group_size": len(group)},
-)
-async def generate_and_rm_group(args, group, sampling_params, evaluation=False):
-    ...
-```
-
-### Choosing a target
-
-`trace_function(...)` needs a trace target, usually a `Sample`, `TraceHandle`, or a list of them.
-
-- Prefer `target="sample"` or `target="group"` when the target is already one of the function arguments.
-- Use `target_getter=...` when the trace target has to be derived from arguments.
-- Avoid relying on automatic inference unless the function signature is simple. The implementation can infer a target from arguments or the current trace context, but explicit targets are more stable and avoid ambiguous traces.
-
-### Recording attrs on decorated functions
-
-If you want to attach attributes at span start, use `attrs_getter=...`:
-
-```python
-@trace_function(
-    "custom_rollout_batch",
-    target="samples",
-    attrs_getter=lambda samples, **_: {"batch_size": len(samples)},
-)
-async def custom_rollout_batch(samples, **kwargs):
-    ...
-```
-
-If you need to add attrs after part of the function has executed, use an inner `trace_span(...)` instead of only relying on the decorator. A common pattern is:
-
-- `trace_function(...)` for the outer function-level lifecycle span
-- nested `trace_span(...)` for important sub-steps such as generation, RM, filtering, or post-processing
-
-If you want to record SGLang generation metadata in a consistent way, reuse `build_sglang_meta_trace_attrs`:
-
-```python
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
-
-with trace_span(sample, "sglang_generate") as span:
-    output = await post(url, payload)
-    span.update(build_sglang_meta_trace_attrs(output["meta_info"]))
-```
-
-## Tips
-
-- Save a small number of rollouts first; the viewer is easiest to read when each dump contains a manageable number of samples.
-- The viewer is built from the saved `.pt` dump, so traces can be inspected offline on another machine.
-- For GPU/kernel-level SGLang profiling traces, see [Profiling](./profiling.md).
-
diff --git a/docs/en/examples/deepseek-r1.md b/docs/en/examples/deepseek-r1.md
index 22e3fd814d..ab10502eec 100644
--- a/docs/en/examples/deepseek-r1.md
+++ b/docs/en/examples/deepseek-r1.md
@@ -171,7 +171,7 @@ OPTIMIZER_ARGS=(
 
 These are the parameters required by sglang. Here, `--rollout-num-gpus-per-engine` basically corresponds to sglang's `tp_size`. Other sglang parameters are passed to slime by adding a `--sglang-` prefix. To fully leverage sglang's large EP inference capabilities, we have added configurations like ep64, dp\_attention dp8, and deepep mode auto.
 
-The final `--rollout-server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128.
+The final `--sglang-server-concurrency` is a parameter specific to slime. It is used to prevent the sglang server's concurrent requests from becoming too large and crashing the HTTP server. The default is 512. However, since we now have one server for 8 nodes, we have adjusted it to 1024 to ensure that each dp rank can have a concurrency of 128.
 
 ```bash
 SGLANG_ARGS=(
@@ -190,7 +190,7 @@ SGLANG_ARGS=(
    --sglang-deepep-mode auto
 
    # make every dp rank have 128 concurrency
-   --rollout-server-concurrency 1024
+   --sglang-server-concurrency 1024
 )
 ```
 
diff --git a/docs/en/examples/glm4.5-355B-A32B.md b/docs/en/examples/glm4.5-355B-A32B.md
new file mode 100644
index 0000000000..1be67a98c1
--- /dev/null
+++ b/docs/en/examples/glm4.5-355B-A32B.md
@@ -0,0 +1,190 @@
+# GLM-4.5 with 64xH100
+
+
+This is an example of doing GLM-4.5 RL training using 64xH100 GPUs.
+
+## Environment Setup
+
+For instructions on setting up the environment and downloading data, please refer to [Example: Qwen3-4B](qwen3-4B.md).
+
+First, you will need to download GLM-4.5 to a directory accessible by all machines (hereinafter referred to as `$BASE_DIR`):
+
+```bash
+hf download zai-org/GLM-4.5 --local-dir $BASE_DIR/GLM-4.5-355B-A32B
+```
+
+Next, we need to convert the huggingface checkpoint into the torch_dist format with 2 nodes, each with 8 GPUs:
+
+```bash
+cd slime/
+source scripts/models/glm4.5-355B-A32B.sh
+PYTHONPATH=/root/Megatron-LM/ torchrun \
+   --nproc-per-node 8 \
+   --master-addr ${MASTER_ADDR} --master-port 12345 \
+   --nnodes=2 --node-rank ${NODE_RANK} \
+   tools/convert_hf_to_torch_dist.py \
+   ${MODEL_ARGS[@]} \
+   --hf-checkpoint $BASE_DIR/GLM-4.5-355B-A32B/ \
+   --save $BASE_DIR/GLM-4.5-355B-A32B_torch_dist/
+```
+
+Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` indicates the node's index, both configured similarly to a multi-node `torchrun` setup.
+
+## Executing the Training
+
+On node0, run:
+
+```bash
+cd slime/
+bash scripts/run-glm4.5-355B-A32B.sh
+```
+
+On other nodes, you need to join the Ray cluster with the following command:
+
+```bash
+ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats"
+```
+
+Alternatively, if you have a list of all node IPs, for example, an MPI hostfile (where each line is `ip slot=8`), you can add the following commands after the `ray start --head` command in `scripts/run-glm4.5-355B-A32B.sh.sh`. This allows you to execute the training entirely from node0:
+
+```bash
+for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do
+  if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" &
+done
+wait
+```
+
+### Parameter Introduction
+
+```bash
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/glm4.5-355B-A32B.sh"
+```
+
+This reads the model's config from [scripts/models/glm4.5-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/models/glm4.5-355B-A32B.sh). These configs are all Megatron parameters. When training with Megatron, it cannot read the model config from the checkpoint, so we need to configure it ourselves. We provide some examples in [scripts/models](https://github.com/THUDM/slime/tree/main/scripts/models/).
+
+#### PERF\_ARGS
+
+A set of Megatron parallelism parameters. Only `--use-dynamic-batch-size` and `--max-tokens-per-gpu` are added by slime.
+
+For the Megatron part, we have configured TP8, PP4, CP2, and EP16.
+
+`max_tokens_per_gpu` refers to the maximum number of tokens each GPU can process. When `use_dynamic_batch_size` is enabled, it will pack data of varying lengths within a batch as close to `max_tokens_per_gpu`. If a single data item exceeds `max_tokens_per_gpu`, it will form its own batch without truncation. When context parallelism (CP) is enabled, it allows CP GPUs to share a total length of `CP * max_tokens_per_gpu` tokens.
+
+When `dynamic_batch_size` is enabled, the traditional `micro_batch_size` is ignored.
+
+⚠️ slime always trains the model using data packing and strictly guarantees per-sample or per-token loss. This means enabling dynamic batch size will not affect the loss calculation. It is recommended to enable it.
+
+```bash
+PERF_ARGS=(
+   --tensor-model-parallel-size 8
+   --sequence-parallel
+   --pipeline-model-parallel-size 4
+   --context-parallel-size 2
+   --expert-model-parallel-size 16
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 16384
+)
+```
+
+#### GRPO\_ARGS
+
+Currently, these are some GRPO-related parameters in slime:
+
+```bash
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+```
+
+If you wish to train without loading the reference model, you need to remove `--use-kl-loss` and set `--kl-coef 0.00` (the default value is 0).
+
+#### OPTIMIZER\_ARGS
+
+We have configured CPU Adam with the following parameters to save GPU memory.
+
+```bash
+OPTIMIZER_ARGS=(
+   ...
+
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+```
+
+#### SGLANG\_ARGS
+
+These are the parameters required by sglang. Here, `--rollout-num-gpus-per-engine` basically corresponds to sglang's `tp_size`. Other sglang parameters are passed to slime by adding a `--sglang-` prefix.
+
+```bash
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 32
+   --sglang-mem-fraction-static 0.7
+   --sglang-enable-dp-attention
+   --sglang-dp-size 4
+)
+```
+
+#### MISC\_ARGS
+
+Some additional Megatron configurations. Note that Megatron's deepep is configured here.
+
+```bash
+MISC_ARGS=(
+   ...
+
+   # use deepep for megatron
+   --moe-enable-deepep
+   --moe-token-dispatcher-type flex
+)
+```
+
+## FP8 Rollout
+
+The open-source FP8 checkpoint of GLM-4.5 is of per-channel quantization, which could not enable deepep in SGLang. We can use the tool scripts within slime to convert an FP8 checkpoint of 128x128 per-block quant:
+
+```bash
+python tools/convert_hf_to_fp8.py \
+    --model-dir $BASE_DIR/GLM-4.5-355B-A32B/ \
+    --save-dir $BASE_DIR/GLM-4.5-355B-A32B-FP8/ \
+    --strategy block --block-size 128 128 \
+    --max-workers 4
+```
+
+Then, simply change `--hf-checkpoint` to `$BASE_DIR/GLM-4.5-355B-A32B-FP8/` to enable FP8 rollout.
+
+And exemplar `SGLANG_ARGS` for FP8 is:
+
+```bash
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 32
+   --sglang-mem-fraction-static 0.7
+   --sglang-enable-dp-attention
+   --sglang-dp-size 32
+   --sglang-ep-size 32
+   --sglang-moe-dense-tp-size 1
+   --sglang-enable-dp-lm-head
+   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 128)
+
+   --sglang-moe-a2a-backend deepep
+   --sglang-deepep-mode auto
+)
+```
diff --git a/docs/en/examples/glm4.7-30B-A3B.md b/docs/en/examples/glm4.7-30B-A3B.md
index e677561aac..1ad877e7e2 100644
--- a/docs/en/examples/glm4.7-30B-A3B.md
+++ b/docs/en/examples/glm4.7-30B-A3B.md
@@ -90,9 +90,9 @@ SGLANG_ARGS=(
    ...
    # MTP speculative decoding (EAGLE)
    --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
+   --sglang-speculative-num-steps 2
    --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --sglang-speculative-num-draft-tokens 3
 )
 ```
 
@@ -102,7 +102,7 @@ This enables SGLang to use the model's MTP layer as a draft model for EAGLE-styl
 
 #### MTP Training
 
-slime also supports training MTP layers jointly with the main model for models that have MTP weight conversion implemented (e.g., MiMo, GLM-4.7). When enabled, the relevant arguments are:
+slime also supports training MTP layers jointly with the main model for models that have MTP weight conversion implemented (e.g., MiMo, GLM-4.5). When enabled, the relevant arguments are:
 
 ```bash
 # Add MTP layer count to model config
@@ -119,9 +119,9 @@ SPEC_ARGS=(
 - `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen.
 - `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2.
 
-> **Note**: MTP training requires the MTP checkpoint bridge to properly convert weights between HuggingFace and Megatron formats. The `GLM4MoELiteBridge` (in `slime_plugins/mbridge/glm4moe_lite.py`) extends the DeepSeek V3 bridge with dynamic MTP layer indexing to support GLM-4.7-Flash's 47-layer architecture.
+> ⚠️ **Note**: MTP training for GLM-4.7-Flash is not yet supported because the deepseek_v3 checkpoint bridge does not include MTP weight conversion (`# TODO: mtp` in upstream mbridge). You can still use MTP for speculative decoding during inference — SGLang handles MTP layers internally.
 >
-> For other models with MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference.
+> For models with full MTP training support (e.g., MiMo), see `scripts/run-mimo-7B-rl-eagle.sh` as a reference.
 
 ### Multi-Node Support
 
diff --git a/docs/en/examples/glm4.7-355B-A32B.md b/docs/en/examples/glm4.7-355B-A32B.md
deleted file mode 100644
index 8035f67824..0000000000
--- a/docs/en/examples/glm4.7-355B-A32B.md
+++ /dev/null
@@ -1,189 +0,0 @@
-# GLM-4.7 with 64xH100
-
-## Environment Preparation
-
-The environment setup and dataset download are the same as for the Qwen3-4B model. You can refer to [Example: Qwen3-4B Model](qwen3-4B.md), replacing mentions of Qwen3-4B with GLM-4.7.
-
-### Prerequisites
-
-GLM-4.7 follows the standard slime Docker environment. For multi-node launches, make sure all nodes can access the same `$BASE_DIR` path and unset proxy variables before starting Ray workers:
-
-```bash
-unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
-```
-
-### Download Model
-
-```bash
-hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B
-```
-
-### Convert Checkpoint
-
-To convert the Hugging Face checkpoint to torch_dist format, use 2 nodes x 8 GPUs:
-
-```bash
-cd /root/slime
-pip install -e . --no-deps
-source scripts/models/glm4.5-355B-A32B.sh
-PYTHONPATH=/root/Megatron-LM/ torchrun \
-   --nproc-per-node 8 \
-   --master-addr ${MASTER_ADDR} --master-port 12345 \
-   --nnodes=2 --node-rank ${NODE_RANK} \
-   tools/convert_hf_to_torch_dist.py \
-   ${MODEL_ARGS[@]} \
-   --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \
-   --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/
-```
-
-Here, `MASTER_ADDR` is the IP of node0, and `NODE_RANK` is the node index, configured just like a multi-node `torchrun` job.
-
-## Run Training
-
-Execute the training script from node0:
-
-```bash
-cd /root/slime
-export BASE_DIR=/shared/path  # accessible by all nodes
-bash scripts/run-glm4.7-355B-A32B.sh
-```
-
-### Parameter Introduction
-
-Here, we briefly introduce the key parts in the [run-glm4.7-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-355B-A32B.sh) script.
-
-#### MoE Configuration
-
-GLM-4.7 is a Mixture-of-Experts (MoE) model with 160 routed experts (top-8 activation) and shared experts. It has 92 layers: 3 dense layers + 89 MoE layers.
-
-1. To support GLM-4.7 on 64xH100, we enable Megatron's CPU Adam to save GPU memory:
-
-   ```bash
-   OPTIMIZER_ARGS=(
-      ...
-      --optimizer-cpu-offload
-      --overlap-cpu-optimizer-d2h-h2d
-      --use-precision-aware-optimizer
-   )
-   ```
-
-2. Enable MoE optimization in Megatron. For the provided 64xH100 example, we use TP=8, PP=4, CP=2, and EP=16:
-
-   ```bash
-   PERF_ARGS=(
-      --tensor-model-parallel-size 8
-      --sequence-parallel
-      --pipeline-model-parallel-size 4
-      --context-parallel-size 2
-      --expert-model-parallel-size 16
-      --expert-tensor-parallel-size 1
-      ...
-      --use-dynamic-batch-size
-      --max-tokens-per-gpu 16384
-   )
-   ```
-
-3. Enable MoE optimization in SGLang with DP attention:
-
-   ```bash
-   SGLANG_ARGS=(
-      --rollout-num-gpus-per-engine 32
-      --sglang-mem-fraction-static 0.7
-      --sglang-enable-dp-attention
-      --sglang-dp-size 4
-      --sglang-ep-size 32
-      --sglang-enable-dp-lm-head
-      --sglang-moe-dense-tp-size 1
-      ...
-   )
-   ```
-
-#### MTP Speculative Decoding (Inference Acceleration)
-
-GLM-4.7 includes MTP (Multi-Token Prediction) layers that can be used for speculative decoding during inference to speed up rollout generation. To enable this, add the following to `SGLANG_ARGS`:
-
-```bash
-SGLANG_ARGS=(
-   ...
-   # MTP speculative decoding (EAGLE)
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
-)
-```
-
-This lets SGLang use the model's MTP layer as the draft model for EAGLE-style speculative decoding.
-
-> ⚠️ **Note**: Speculative decoding requires additional GPU memory. If you encounter OOM issues, try reducing `--sglang-mem-fraction-static` or disabling speculative decoding.
-
-#### MTP Training
-
-slime also supports training the MTP layers jointly with the main model for GLM-4.7. When enabled, the relevant arguments are:
-
-```bash
-# Add MTP layer count to model config
-MODEL_ARGS+=(--mtp-num-layers 1)
-
-# Enable MTP training
-MTP_ARGS=(
-   --enable-mtp-training
-   --mtp-loss-scaling-factor 0.2
-)
-```
-
-- `--mtp-num-layers 1`: Tells Megatron to load the MTP layer from the checkpoint.
-- `--enable-mtp-training`: Enables gradient computation for MTP layers. Without this flag, the MTP layer is loaded but frozen.
-- `--mtp-loss-scaling-factor 0.2`: Weight of the MTP loss relative to the main policy loss. Default is 0.2.
-
-> **Note**: MTP training for GLM-4.7 relies on `GLM4MoEBridge` (in `slime_plugins/mbridge/glm4moe.py`) to map regular and MTP weights between HuggingFace and Megatron formats.
-
-#### Multi-Node Support
-
-This example already targets multi-node training. Before launching:
-
-- Place the model checkpoints and datasets on a path accessible by all nodes.
-- Set `MASTER_ADDR` to an address reachable by all nodes.
-- Unset proxy variables before starting Ray workers.
-- Provide a `HOSTFILE` listing worker IPs (one per line) and export `HOSTFILE=/path/to/hostfile` before launching.
-- Adjust parallelism coherently. The default example uses TP=8, PP=4, EP=16, CP=2, while rollout uses 32 GPUs per engine with SGLang DP attention.
-
-If your rollout GPU count does not divide the expert count cleanly, you can use `--sglang-ep-num-redundant-experts` to add redundant experts.
-
-## FP8 Rollout
-
-The open-source FP8 checkpoint of GLM-4.7 uses per-channel quantization, which cannot currently enable DeepEP in SGLang. You can convert it to a 128x128 per-block FP8 checkpoint with the tool provided in slime:
-
-```bash
-cd /root/slime
-python tools/convert_hf_to_fp8.py \
-    --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \
-    --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \
-    --strategy block --block-size 128 128 \
-    --max-workers 4
-```
-
-Then switch `--hf-checkpoint` to `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` to enable FP8 rollout.
-
-An example FP8 `SGLANG_ARGS` setup is:
-
-```bash
-SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 32
-   --sglang-mem-fraction-static 0.7
-   --sglang-enable-dp-attention
-   --sglang-dp-size 32
-   --sglang-ep-size 32
-   --sglang-moe-dense-tp-size 1
-   --sglang-enable-dp-lm-head
-   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 128)
-
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
-
-   --sglang-moe-a2a-backend deepep
-   --sglang-deepep-mode auto
-)
-```
diff --git a/docs/en/examples/qwen3-4B.md b/docs/en/examples/qwen3-4B.md
index 688b30b03c..56711878f6 100644
--- a/docs/en/examples/qwen3-4B.md
+++ b/docs/en/examples/qwen3-4B.md
@@ -280,10 +280,10 @@ In this case, 2 GPUs will be allocated for training, and 6 GPUs will be allocate
 
 ⚠️  If the concurrency on each sglang server is too high, it may exceed sglang's default CUDA graph concurrency limit (the default maximum is 160), which will affect inference speed. You can adjust this in the following two ways:
 
-1.  Use `--rollout-server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example:
+1.  Use `--sglang-server-concurrency` to limit the maximum number of concurrent requests sent to a single sglang server. For example:
 
     ```bash
-    --rollout-server-concurrency 160
+    --sglang-server-concurrency 160
     ```
 
 2.  Use `--sglang-cuda-graph-bs` (which corresponds to sglang's native `--cuda-graph-bs` argument) to increase the number of CUDA graphs initialized by sglang. For example:
diff --git a/docs/en/get_started/customization.md b/docs/en/get_started/customization.md
index 79d2e8264f..bcbd993ef2 100644
--- a/docs/en/get_started/customization.md
+++ b/docs/en/get_started/customization.md
@@ -28,6 +28,7 @@ Below is a summary of all available customization interfaces and their purposes.
 | [`--custom-megatron-init-path`](#17-megatron-hooks) | Custom initialization after Megatron setup. |
 | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hooks) | Custom logic before log probability computation. |
 | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hooks) | Custom logic before each training step. |
+| [`--slime-router-middleware-paths`](#18-slime-router-middleware---slime-router-middleware-paths) | Add custom middleware to slime router. |
 
 ## Detailed Interface Reference
 
@@ -399,14 +400,27 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler
 
 ---
 
-### 18. MoE Routing Replay
+### 18. slime Router Middleware (`--slime-router-middleware-paths`)
+
+**Purpose**: Add custom middleware to the slime router for request processing.
+
+**Use Cases**:
+- Request/response transformation
+- Custom routing logic
+- Caching and optimization
+
+---
+
+### 19. MoE Routing Replay
 
 Stabilize MoE RL training by recording and replaying expert routing decisions to ensure consistency.
 
 | Argument | Description |
 | --- | --- |
 | `--use-routing-replay` | Forward-backward routing consistency in training. ([arXiv:2507.18071](https://arxiv.org/abs/2507.18071)) |
-| `--use-rollout-routing-replay` | R3: Replay routing from rollout during training. Supported by slime's default `sglang_router` path. ([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) |
+| `--use-rollout-routing-replay` | R3: Replay routing from rollout during training. **Requires `--use-slime-router`**. ([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) |
+
+For detailed explanation of R3 and SlimeRouter, see [Slime Router](../advanced/slime-router.md).
 
 ## Testing Custom Function Paths
 
diff --git a/docs/en/get_started/quick_start.md b/docs/en/get_started/quick_start.md
index 5d3856fdeb..3d86b17141 100644
--- a/docs/en/get_started/quick_start.md
+++ b/docs/en/get_started/quick_start.md
@@ -71,7 +71,7 @@ hf download --repo-type dataset zhuzilin/aime-2024 \
 
 When using Megatron as the training backend, you need to first convert Hugging Face format model weights to Megatron `torch_dist` format.
 
-First, load the configuration file of the target model. The `slime/scripts/models` directory contains configuration files for supported models. You need to `source` the corresponding model script to load the configuration parameters into the current environment. Here we use GLM4-9B model as an example, and it's similar for Qwen3-4B, Qwen3.5, Qwen3.6, GLM-4.7-Flash, Qwen3-30B-A3B, etc.
+First, load the configuration file of the target model. The `slime/scripts/models` directory contains configuration files for supported models. You need to `source` the corresponding model script to load the configuration parameters into the current environment. Here we use GLM4-9B model as an example, and it's similar for Qwen3-4B, GLM-4.7-Flash, Qwen3-30B-A3B, etc.
 
 ```bash
 cd /root/slime
@@ -581,6 +581,6 @@ export NVSHMEM_BOOTSTRAP_UID_SOCK_IFNAME=$(ip -o -4 addr show | awk '$4 ~ /^10\.
 slime has been deeply optimized for distributed training of large-scale Mixture of Experts (MoE) models. We provide some end-to-end training cases for reference:
 
 - [Example: 8xH100 Training GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md)
-- [Example: 64xH100 Training GLM-4.7](../examples/glm4.7-355B-A32B.md)
+- [Example: 64xH100 Training GLM-4.5](../examples/glm4.5-355B-A32B.md)
 - [Example: 128xH100 Training DeepSeek-R1](../examples/deepseek-r1.md)
 - The scripts such as `scripts/run_qwen3_30b_a3b.py`, `scripts/run_glm45_355b_a32b.py` also support multi-node training, though there are little documentations about it currently.
diff --git a/docs/en/get_started/usage.md b/docs/en/get_started/usage.md
index cbd0122208..b85e537028 100644
--- a/docs/en/get_started/usage.md
+++ b/docs/en/get_started/usage.md
@@ -35,6 +35,7 @@ Additionally, slime supports Prefill and Decode disaggregation (PD Disaggregatio
 slime supports multiple training backends, which can be selected via the `--train-backend` parameter:
 
 - `megatron` (default): Uses Megatron-LM as the training backend, supporting efficient training of large-scale models.
+- `fsdp` (experimental): Uses PyTorch FSDP as the training backend, allowing direct loading of HuggingFace format weights without conversion.
 
 ### Loading Megatron
 
@@ -189,6 +190,7 @@ Additionally, we provide a `metadata_key`, which defaults to `"metadata"`. When
   Note: On-policy distillation (OPD) is now orthogonal to the advantage estimator. Use `--use-opd` and `--opd-kl-coef` to enable OPD on top of any estimator.
 - `--calculate-per-token-loss`: By default, slime calculates loss on a per-sample basis, i.e., `mean(sum(sample_i) / len(sample_i))`. Enable this flag to calculate loss on a per-token basis, i.e., `sum(sum(sample_i)) / sum(len(sample_i))`.
 - `--use-tis`: Enable this setting to use TIS (Truncated Importance Sampling) (https://fengyao.notion.site/off-policy-rl).
+- `--true-on-policy-mode`: Enable True On-Policy mode, which strictly ensures that data is generated by the current policy during training.
 
 #### GRPO Algorithm
 
@@ -256,24 +258,6 @@ PPO-related parameters:
 - `--value-clip`: Clip range for value loss.
 - `--kl-coef`: KL penalty coefficient for reward shaping.
 
-### Advanced Megatron Configuration (--megatron-config-path)
-
-For PPO workflows, you can use `--megatron-config-path` with a YAML file to override Megatron arguments separately for actor and critic. Common use cases include setting a different critic `lr`, or giving actor and critic different `load` / `save` paths.
-
-```yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-> **Note:** This configuration currently only supports PPO, and in current PPO the actor and critic must use the same Megatron parallel topology. The recommended pattern is to keep parallelism-related settings in the shared CLI arguments and put only role-specific differences in YAML. See [Megatron Config: Role-Based Training Overrides](../advanced/megatron-config.md) for details.
-
 ## Custom Rollout Function
 
 slime supports customizing data generation (rollout) to various degrees.
@@ -382,36 +366,6 @@ After starting, all SGLang servers will register with the router via the `/add_w
 
 When you configure an external router using `--sglang-router-ip` and `--sglang-router-port`, slime will not start an internal router. Instead, it will register all its servers with this external router. You can then use this external router's address to implement more complex data generation workflows. Note that the router supports OpenAI-compatible APIs.
 
-### Advanced Engine Configuration (--sglang-config)
-
-For advanced deployments, you can use `--sglang-config` with a YAML file to configure server groups, multi-model serving, and selective weight updates.
-
-**Multi-model deployment** allows serving multiple models simultaneously (e.g., an actor model that receives weight updates and a frozen reference/reward model):
-
-```yaml
-sglang:
-  - name: actor
-    update_weights: true          # receives weight updates from training (default)
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-  - name: ref
-    model_path: /path/to/ref_model
-    update_weights: false          # frozen, no weight updates
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-Each model gets its own router. The per-model router info is accessible via `args.sglang_model_routers` (a dict mapping model name to `(ip, port)` tuples). Custom rollout functions can use `get_model_url(args, "ref")` from `slime.rollout.sglang_rollout` to route requests to a specific model.
-
-**Server group features:**
-- `worker_type`: `regular`, `prefill`, `decode`, or `placeholder` (reserves GPU slots without creating engines)
-- `overrides`: Dict of SGLang `ServerArgs` field overrides applied on top of `--sglang-*` CLI args
-- `num_gpus_per_engine`: Per-group TP size override
-
 ## How to Use Megatron
 
 slime supports different and lightly modified versions of Megatron by reusing common functions from the `megatron.training` directory, such as `parse_args`, `save_checkpoint`, and `load_checkpoint`. Therefore, when using it, you must ensure that Megatron is accessible in the `PYTHONPATH`, for example, by adding `export PYTHONPATH=/root/Megatron-LM` at runtime.
diff --git a/docs/en/index.rst b/docs/en/index.rst
index d5ca098e06..162ba19baf 100644
--- a/docs/en/index.rst
+++ b/docs/en/index.rst
@@ -34,21 +34,20 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models
 
    examples/glm4.7-30B-A3B.md
    examples/qwen3-30B-A3B.md
-   examples/glm4.7-355B-A32B.md
+   examples/glm4.5-355B-A32B.md
    examples/deepseek-r1.md
 
 .. toctree::
    :maxdepth: 1
    :caption: Advanced Features
 
+   advanced/slime-router.md
    advanced/on-policy-distillation.md
    advanced/speculative-decoding.md
    advanced/low-precision.md
    advanced/reproducibility.md
    advanced/fault-tolerance.md
    advanced/pd-disaggregation.md
-   advanced/sglang-config.md
-   advanced/megatron-config.md
    advanced/arch-support-beyond-megatron.md
 
 .. toctree::
@@ -67,7 +66,6 @@ slime is the RL-framework behind GLM-4.7, GLM-4.6 and GLM-4.5. Apart from models
 
    developer_guide/ci.md
    developer_guide/debug.md
-   developer_guide/trace.md
    developer_guide/profiling.md
 
 .. toctree::
diff --git a/docs/en/vllm/ROUTER_DESIGN.md b/docs/en/vllm/ROUTER_DESIGN.md
deleted file mode 100644
index c18a29c997..0000000000
--- a/docs/en/vllm/ROUTER_DESIGN.md
+++ /dev/null
@@ -1,462 +0,0 @@
-# RFC: Replace SGLang Backend with vLLM — Router Integration
-
----
-
-## Summary
-
-Replace the SGLang inference backend behind **SlimeRouter** with **vLLM** while keeping the existing router and middleware stack completely unchanged.
-This RFC covers **only the router layer** — what APIs the vLLM backend must expose, how the existing SlimeRouter is reused, and what translation is needed between the two formats.
-
-**Key design decision:** Reuse vLLM's built-in [OpenAI-compatible API server](https://docs.vllm.ai/en/stable/serving/openai_compatible_server/) (`vllm serve`) 
-
-
----
-
-## 1. Target Architecture
-
-```
- Rollout Workers                    SlimeRouter (NO CHANGE)                vLLM Engines (NEW)
- ──────────────                    ────────────────────────                ──────────────────
-                                   ┌──────────────────────┐
- POST /generate ──────────────────▶│ RadixTreeMiddleware   │
-                                   │  • prefix cache       │
-                                   │  • retry on abort     │
-                                   │  • token/logprob cache│
-                                   └──────────┬───────────┘
-                                              │
-                                   ┌──────────▼───────────┐
-                                   │ SlimeRouter.proxy()   │         ┌─────────────────────┐
-                                   │  • least-connections  │────────▶│ vLLM Translation    │
-                                   │    load balancer      │         │ Sidecar (per engine) │
-                                   │  • health check loop  │         │                     │
-                                   └──────────────────────┘         │ POST /generate      │
-                                                                     │   ↓ translate        │
-                                                                     │ POST /v1/completions │
-                                                                     │   ↓ translate back   │
-                                                                     │ → SGLang-format JSON │
-                                                                     └─────────┬───────────┘
-                                                                               │
-                                                                     ┌─────────▼───────────┐
-                                                                     │ vLLM Server          │
-                                                                     │ (vllm serve)         │
-                                                                     │  • /v1/completions   │
-                                                                     │  • /health           │
-                                                                     │  • /sleep, /wake_up  │
-                                                                     │  • /pause, /resume   │
-                                                                     │  • /update_weights   │
-                                                                     └─────────────────────┘
-```
-
-### What stays the same
-
-| Component | Change | Reason |
-|---|---|---|
-| `SlimeRouter` ([router.py](slime/router/router.py)) | **None** | Engine-agnostic HTTP proxy; only reads JSON responses |
-| `RadixTreeMiddleware` ([radix_tree_middleware.py](slime/router/middleware_hub/radix_tree_middleware.py)) | **None** | Operates on request/response JSON; has no engine-specific code |
-| `StringRadixTrie` ([radix_tree.py](slime/router/middleware_hub/radix_tree.py)) | **None** | Pure data structure, no engine coupling |
-| Middleware loading (`--slime-router-middleware-paths`) | **None** | Dynamic import via `load_function()` |
-
-### What is new
-
-| Component | Description |
-|---|---|
-| `vllm_translation_sidecar.py` | Lightweight FastAPI process co-located with each vLLM engine. Receives SGLang-format `/generate` requests, translates to vLLM's `/v1/completions`, translates responses back. Also proxies lifecycle endpoints (`/abort_request`, `/health_generate`, etc.). |
-| `vllm_engine.py` | Ray actor that manages the vLLM server process lifecycle (via `vllm serve`), the translation sidecar, weight updates, and registration with the router. |
-
----
-
-## 2. Reusing SlimeRouter — Zero Modification
-
-The SlimeRouter communicates with backends through **five interaction points**. All are already engine-agnostic:
-
-### 2.1 Worker Registration
-
-**Flow:** Engine starts → engine calls `POST /add_worker?url=http://{host}:{port}` → router adds to pool.
-
-```
-Router state after registration:
-  worker_request_counts["http://10.0.0.1:10090"] = 0
-  worker_failure_counts["http://10.0.0.1:10090"] = 0
-```
-
-**vLLM action:** The `VLLMEngine` Ray actor calls this endpoint after verifying the vLLM server + translation sidecar are healthy. The registered URL points to the **sidecar**, not the raw vLLM server. No router change needed.
-
-### 2.2 Request Proxying
-
-**Flow:** `POST /generate` → middleware pipeline → `SlimeRouter.proxy()` → `httpx` forwards to backend (sidecar).
-
-The router selects a backend via **least-connections** (`_use_url()`), forwards the raw request body as-is, and returns the response as-is. It never inspects or transforms the request/response payload.
-
-**vLLM action:** The sidecar receives the forwarded request, translates it to `/v1/completions`, calls the co-located vLLM server, translates the response back to SGLang format, and returns it.
-
-### 2.3 Health Check
-
-**Flow:** Background loop calls `GET {worker_url}/health` every N seconds.
-
-- 200 → healthy, reset failure count
-- Non-200 or timeout → increment failure count
-- Failures ≥ threshold (default 3) → quarantine worker permanently
-
-**vLLM action:** The sidecar's `/health` proxies to vLLM's built-in `/health` endpoint (returns 200 when ready). Compatible out of the box.
-
-### 2.4 Worker Listing
-
-**Flow:** `GET /list_workers` → returns `{"urls": [...]}`
-
-Used by the rollout to discover engines for direct abort calls. No engine involvement.
-
-### 2.5 Retrieve from Text (Radix Tree)
-
-**Flow:** `POST /retrieve_from_text` → router looks up the radix tree cache → returns tokens/logprobs.
-
-Fully router-internal. Never reaches the engine.
-
----
-
-## 3. API Contract — What the Translation Sidecar Must Expose
-
-The translation sidecar sits between SlimeRouter and the vLLM server. It receives SGLang-format requests and returns SGLang-format responses.
-
-### 3.1 `POST /generate` — Generation
-
-This is the primary endpoint. The sidecar translates between Slime's format and vLLM's `/v1/completions`.
-
-#### Incoming Request (from router)
-
-```json
-{
-  "input_ids": [128000, 2610, 553, 264, 11190, 18328, 13],
-  "input_tokens": [128000, 2610, 553, 264, 11190, 18328, 13],
-  "sampling_params": {
-    "temperature": 0.7,
-    "top_p": 0.9,
-    "top_k": -1,
-    "max_new_tokens": 1024,
-    "stop": ["<|endoftext|>"],
-    "stop_token_ids": [128001],
-    "skip_special_tokens": false,
-    "no_stop_trim": true,
-    "spaces_between_special_tokens": false
-  },
-  "return_logprob": true,
-  "stream": false
-}
-```
-
-#### Translated Request (to vLLM `/v1/completions`)
-
-```json
-{
-  "model": "<model_name>",
-  "prompt": [128000, 2610, 553, 264, 11190, 18328, 13],
-  "max_tokens": 1024,
-  "temperature": 0.7,
-  "top_p": 0.9,
-  "top_k": -1,
-  "stop": ["<|endoftext|>"],
-  "stop_token_ids": [128001],
-  "skip_special_tokens": false,
-  "include_stop_str_in_output": true,
-  "spaces_between_special_tokens": false,
-  "logprobs": 1,
-  "stream": false,
-  "extra_body": {
-    "return_token_ids": true
-  }
-}
-```
-
-**Key translations:**
-- `input_ids` → `prompt` (vLLM accepts `list[int]` as pre-tokenized prompt)
-- `max_new_tokens` → `max_tokens`
-- `no_stop_trim: true` → `include_stop_str_in_output: true`
-- `return_logprob: true` → `logprobs: 1` + `extra_body.return_token_ids: true`
-
-#### vLLM Response (from `/v1/completions`)
-
-```json
-{
-  "id": "cmpl-abc123",
-  "choices": [{
-    "text": "I'll help you with that. The answer is 42.",
-    "logprobs": {
-      "token_logprobs": [-0.152, -0.089, -0.203],
-      "tokens": ["I", "'ll", " help"]
-    },
-    "token_ids": [40, 3358, 1520],
-    "finish_reason": "stop"
-  }],
-  "usage": {
-    "prompt_tokens": 7,
-    "completion_tokens": 3,
-    "total_tokens": 10
-  }
-}
-```
-
-#### Translated Response (returned to router)
-
-```json
-{
-  "text": "I'll help you with that. The answer is 42.",
-  "output_ids": [40, 3358, 1520],
-  "meta_info": {
-    "output_token_logprobs": [
-      [-0.152, 40],
-      [-0.089, 3358],
-      [-0.203, 1520]
-    ],
-    "finish_reason": {
-      "type": "stop"
-    },
-    "weight_version": 3,
-    "prompt_tokens": 7,
-    "cached_tokens": 0
-  }
-}
-```
-
-##### Field-by-field contract
-
-| Field | Type | Required | Consumer | Description |
-|---|---|---|---|---|
-| `text` | `str` | **Yes** | Rollout, Middleware | Generated text (output only, not including prompt) |
-| `output_ids` | `list[int]` | **Yes** | Middleware | Generated token IDs. Middleware checks existence as a gate for caching. |
-| `meta_info.output_token_logprobs` | `list[[float, int]]` | **Yes** (if `return_logprob`) | Rollout, Middleware | Each element is `[logprob, token_id]`. Used for RL policy ratio calculation. |
-| `meta_info.finish_reason` | `{"type": str}` | **Yes** | Rollout, Middleware | Must be `{"type": "stop"}`, `{"type": "length"}`, or `{"type": "abort"}`. **Not** a plain string. |
-| `meta_info.weight_version` | `int` | **Yes** | Middleware, Rollout | Current model weight version. Tracked by the sidecar (incremented on each weight update). |
-| `meta_info.prompt_tokens` | `int` | Nice-to-have | Rollout (stats) | From `usage.prompt_tokens`. |
-| `meta_info.cached_tokens` | `int` | Nice-to-have | Rollout (stats) | vLLM doesn't expose this directly; default to `0`. |
-
-### 3.2 `GET /health` — Health Check
-
-```
-GET /health
-→ Sidecar proxies to vLLM's GET /health
-→ 200 OK        (engine ready)
-→ 503 or timeout (engine not ready / overloaded)
-```
-
-vLLM already provides this endpoint. **Passthrough — no translation needed.**
-
-### 3.3 `POST /abort_request` — Cancel Generation
-
-```
-POST /abort_request
-Body: {"abort_all": true}
-→ 200 OK
-```
-
-Called **directly** by the rollout to each engine (bypasses the router). The rollout discovers engine URLs via `GET /list_workers`, then sends abort to each.
-
-**vLLM approach:** vLLM uses **HTTP connection close** for abort (via its `@with_cancellation` decorator). When a client disconnects, the in-flight request is automatically cancelled.
-
-**Implementation options:**
-1. **Track active connections.** The sidecar maintains a set of active `httpx` connections to the vLLM server. On `POST /abort_request`, close all of them — triggering vLLM's cancellation.
-2. **Use vLLM's `/pause` endpoint.** Call `POST /pause` to block new requests, then `POST /resume` after the RL training step completes. This is semantically closer to how Slime uses abort (clearing the decks between training generations).
-
-> **Note:** vLLM has `POST /abort_requests` only in disaggregated mode. For standard mode, HTTP disconnect is the canonical abort mechanism.
-
-### 3.4 `GET /health_generate` — Startup Readiness Probe
-
-```
-GET /health_generate
-→ 200 OK        (model loaded, engine ready for generation)
-```
-
-Called by `VLLMEngine.init()` during startup to block until the engine is fully ready. The sidecar implements this by calling vLLM's `GET /health` and optionally performing a dummy `/v1/completions` call with `max_tokens=1` to verify end-to-end readiness.
-
-### 3.5 Sampling Params Translation
-
-The request uses SGLang-format parameter names. The sidecar translates to vLLM's `/v1/completions` format:
-
-| SGLang field (in request) | vLLM `/v1/completions` field | Notes |
-|---|---|---|
-| `input_ids` | `prompt` | Direct — vLLM accepts `list[int]` as pre-tokenized prompt |
-| `temperature` | `temperature` | Direct |
-| `top_p` | `top_p` | Direct |
-| `top_k` | `top_k` | Both use `-1` for disabled |
-| `max_new_tokens` | `max_tokens` | **Name change** |
-| `stop` | `stop` | Direct (list of strings) |
-| `stop_token_ids` | `stop_token_ids` | Direct |
-| `skip_special_tokens` | `skip_special_tokens` | Direct |
-| `no_stop_trim` | `include_stop_str_in_output` | **Same semantics, different name** |
-| `spaces_between_special_tokens` | `spaces_between_special_tokens` | Direct |
-| `return_logprob` | `logprobs` (set to `1`) | Also add `extra_body.return_token_ids = true` |
-| `sampling_seed` | `seed` | Optional |
-| — | `model` | Must be set to the model name served by vLLM |
-
-### 3.6 Response Translation Pseudocode
-
-```python
-def translate_vllm_response(vllm_resp: dict, weight_version: int) -> dict:
-    """Translate vLLM /v1/completions response to SGLang format."""
-    choice = vllm_resp["choices"][0]
-    usage = vllm_resp.get("usage", {})
-
-    # Build output_token_logprobs: zip logprobs with token IDs
-    output_token_logprobs = None
-    if choice.get("logprobs") and choice.get("token_ids"):
-        output_token_logprobs = [
-            [logprob, token_id]
-            for logprob, token_id in zip(
-                choice["logprobs"]["token_logprobs"],
-                choice["token_ids"]
-            )
-        ]
-
-    # Translate finish_reason: plain string → {"type": str}
-    raw_reason = choice.get("finish_reason")
-    finish_reason = {"type": raw_reason if raw_reason else "abort"}
-
-    return {
-        "text": choice["text"],
-        "output_ids": choice.get("token_ids", []),
-        "meta_info": {
-            "output_token_logprobs": output_token_logprobs,
-            "finish_reason": finish_reason,
-            "weight_version": weight_version,
-            "prompt_tokens": usage.get("prompt_tokens", 0),
-            "cached_tokens": 0,
-        }
-    }
-```
-
-### 3.7 `finish_reason` Translation Table
-
-| vLLM returns | Translate to | Notes |
-|---|---|---|
-| `"stop"` | `{"type": "stop"}` | Normal completion |
-| `"length"` | `{"type": "length"}` | Hit `max_tokens` |
-| `None` (aborted/incomplete) | `{"type": "abort"}` | Triggers middleware retry logic (sleep 30s, up to 5 retries) |
-
----
-
-## 4. Server Launch Configuration
-
-The `VLLMEngine` Ray actor should launch vLLM as follows:
-
-```bash
-# Environment
-export VLLM_SERVER_DEV_MODE=1
-
-# Launch vLLM server
-vllm serve <model_path> \
-    --host 0.0.0.0 \
-    --port <engine_port> \
-    --tensor-parallel-size <tp_size> \
-    --enable-sleep-mode \
-    --enforce-eager \
-    --gpu-memory-utilization 0.9 \
-    --disable-log-requests
-```
-
-The translation sidecar runs on a separate port (`<sidecar_port>`) and is the URL registered with the router via `POST /add_worker?url=http://{host}:{sidecar_port}`.
-
-```
-                Router
-                  │
-                  ▼
-    ┌─────────────────────────┐
-    │ Translation Sidecar     │  ◄── registered with router
-    │ port: sidecar_port      │
-    │                         │
-    │ /generate ──translate──▶│──┐
-    │ /health ──passthrough──▶│  │
-    │ /abort_request          │  │
-    │ /health_generate        │  │
-    └─────────────────────────┘  │
-                                 │
-    ┌─────────────────────────┐  │
-    │ vLLM Server             │◄─┘
-    │ port: engine_port       │
-    │                         │
-    │ /v1/completions         │
-    │ /health                 │
-    │ /sleep, /wake_up        │
-    │ /pause, /resume         │
-    │ /update_weights         │
-    │ /init_weight_transfer   │
-    └─────────────────────────┘
-```
-
----
-
-## 5. Abort Strategy — Detailed Design
-
-vLLM's abort mechanism differs fundamentally from SGLang's:
-
-| Aspect | SGLang | vLLM |
-|---|---|---|
-| Abort granularity | Per-request via `POST /abort_request` with `rid` | Per-connection via HTTP disconnect |
-| Bulk abort | `{"abort_all": true}` | No built-in equivalent |
-| Mechanism | Engine tracks `request_id`, explicit `abort()` | `@with_cancellation` decorator; request cancelled when client disconnects |
-| Between-generation abort | Abort + restart | `POST /pause` → training → `POST /resume` |
-
-### Recommended implementation
-
-For the Slime RL use case, the rollout calls `abort_all` between generation rounds (to clear the engine before the next batch). The best vLLM equivalent is:
-
-```python
-# In the translation sidecar
-@app.post("/abort_request")
-async def abort_request(request: Request):
-    body = await request.json()
-    if body.get("abort_all"):
-        # Option 1: Close all tracked httpx connections → triggers vLLM cancellation
-        for conn in active_connections:
-            await conn.aclose()
-        active_connections.clear()
-
-        # Option 2: Use pause/resume (cleaner)
-        await httpx.post(f"{vllm_url}/pause")
-        await httpx.post(f"{vllm_url}/resume")
-
-    return {"status": "ok"}
-```
-
----
-
-## 6. Endpoints Summary — Gap Analysis
-
-### Engine-side endpoints (vLLM built-in vs. needs implementation)
-
-| Endpoint | SGLang | vLLM Built-in | Action |
-|---|---|---|---|
-| `POST /v1/completions` | — | ✅ | **Reuse** — target for translation |
-| `GET /health` | ✅ | ✅ | **Reuse** as-is (passthrough) |
-| `POST /pause` | — | ✅ (dev mode) | **Reuse** for abort/weight-update |
-| `POST /resume` | — | ✅ (dev mode) | **Reuse** for abort/weight-update |
-| `POST /sleep` | — | ✅ (dev mode) | **Reuse** for weight updates |
-| `POST /wake_up` | — | ✅ (dev mode) | **Reuse** for weight updates |
-| `POST /collective_rpc` | — | ✅ (dev mode) | **Reuse** for weight reload |
-| `GET /is_sleeping` | — | ✅ (dev mode) | **Reuse** for state checks |
-| `POST /init_weight_transfer_engine` | — | ✅ (dev mode) | **Reuse** for NCCL setup |
-| `POST /update_weights` | — | ✅ (dev mode) | **Reuse** for NCCL weight apply |
-| `GET /get_world_size` | — | ✅ (dev mode) | **Reuse** for TP world size |
-
-### Translation sidecar endpoints (to implement)
-
-| Endpoint | Description | Complexity |
-|---|---|---|
-| `POST /generate` | Translate SGLang → `/v1/completions` → SGLang | **Medium** — main logic |
-| `GET /health` | Proxy to vLLM `/health` | **Trivial** |
-| `GET /health_generate` | Health + optional dummy completion | **Low** |
-| `POST /abort_request` | Close connections or pause/resume | **Low** |
-| `GET /flush_cache` | `POST /sleep?level=1` + `POST /wake_up?tags=kv_cache` | **Low** |
-| `GET /get_weight_version` | Return sidecar-tracked version counter | **Trivial** |
-
-### Router endpoints (no change needed)
-
-| Endpoint | Action |
-|---|---|
-| `POST /add_worker` | No change |
-| `GET /list_workers` | No change |
-| `POST /retrieve_from_text` | No change |
-| Catch-all proxy | No change |
-
----
-
-
-
-
diff --git a/docs/zh/advanced/megatron-config.md b/docs/zh/advanced/megatron-config.md
deleted file mode 100644
index 1616c9f830..0000000000
--- a/docs/zh/advanced/megatron-config.md
+++ /dev/null
@@ -1,131 +0,0 @@
-# Megatron Config：按角色覆盖训练参数
-
-`--megatron-config-path` 是一个基于 YAML 的配置系统，用于在公共 Megatron CLI 参数之上，为不同训练角色追加覆盖。目前它主要用于 PPO 场景中的 actor / critic 配置。
-
-与 `--sglang-config` 不同，`--megatron-config-path` 不负责部署、路由或 GPU 资源编排；它只负责决定“这个角色最终使用哪些训练参数”。
-
----
-
-## 设计概览
-
-默认情况下（不使用 `--megatron-config-path`），actor 和 critic 都直接继承命令行中的 Megatron / slime 参数。
-
-使用 `--megatron-config-path` 后，配置会分成两层：
-
-- **公共 CLI 参数**：定义共享的 Megatron 拓扑、资源申请和默认训练参数。
-- **角色 YAML 覆盖**：只覆盖 actor / critic 之间真正不同的配置项。
-
-**核心设计原则：**
-
-- **CLI 是公共基线。** slime 会先解析普通命令行参数，再应用 YAML 中的角色覆盖。
-- **缺失角色自动继承。** 如果某个 role 没有出现在 YAML 中，该角色会直接继承 CLI 参数。
-- **资源申请仍由 CLI 控制。** YAML 中的 `num_nodes` 和 `num_gpus_per_node` 会被忽略；资源分配仍由 `--actor-num-*` / `--critic-num-*` 决定。
-
----
-
-## 配置格式
-
-配置文件是一个 YAML 文档，顶层 `megatron` 键包含一个角色定义列表：
-
-```yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-      save: /path/to/actor_ckpt
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-      save: /path/to/critic_ckpt
-```
-
-### 字段说明
-
-| 字段 | 类型 | 默认值 | 说明 |
-|------|------|--------|------|
-| `name` | `str` | 可选 | 配置项标签。当前运行时不会依赖它，建议保留为 `default`，方便未来扩展。 |
-| `role` | `str` | **必填** | 角色名。目前支持 `actor` 和 `critic`。 |
-| `overrides` | `dict` | `{}` | 该角色的参数覆盖。会在公共 CLI 参数之上应用。 |
-| `args` | `dict` | `{}` | `overrides` 的向后兼容别名。新配置建议统一使用 `overrides`。 |
-
-> **注意：** `overrides` 中的 key 使用 argparse 属性名，而不是命令行 flag 名。例如写 `tensor_model_parallel_size`，而不是 `tensor-model-parallel-size`。
-
----
-
-## 使用方式
-
-一个典型的 PPO 用法如下：
-
-```yaml
-# megatron_ppo.yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-```bash
-python train.py \
-  --advantage-estimator ppo \
-  --use-critic \
-  --megatron-config-path megatron_ppo.yaml \
-  --tensor-model-parallel-size 2 \
-  --sequence-parallel \
-  --pipeline-model-parallel-size 1 \
-  --context-parallel-size 1 \
-  --expert-model-parallel-size 1 \
-  --expert-tensor-parallel-size 1 \
-  --actor-num-nodes 1 \
-  --actor-num-gpus-per-node 8 \
-  --critic-num-nodes 1 \
-  --critic-num-gpus-per-node 8 \
-  ...
-```
-
-在这个模式下：
-
-- CLI 负责共享的并行策略和资源配置；
-- YAML 负责 actor / critic 的差异项，比如 `lr`、`load`、`save`、optimizer 或 scheduler 相关参数。
-
-### 只覆盖一个角色
-
-你也可以只为一个角色写覆盖，另一个角色自动继承公共 CLI 参数。例如只调整 critic 学习率：
-
-```yaml
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-这时 actor 会继续直接使用命令行中的公共参数。
-
----
-
-## 当前限制
-
-- **目前只支持 PPO。** `--megatron-config-path` 当前主要用于 PPO 工作流中的 actor / critic 角色配置。对于 GRPO、REINFORCE++ 等不依赖 critic 的流程，目前不建议使用这套角色配置。
-- **当前 PPO 下，actor 和 critic 的 Megatron 并行配置必须一致。** 特别是 `tensor_model_parallel_size`、`pipeline_model_parallel_size`、`context_parallel_size`、`expert_model_parallel_size`、`expert_tensor_parallel_size`、`sequence_parallel` 等拓扑相关参数，不应在 actor 和 critic 之间配置成不同的值。
-- **推荐把并行相关参数继续放在 CLI 中。** 当前最稳妥的用法是：并行与资源参数写在公共 CLI 中，只在 YAML 中覆盖角色差异项，例如 `lr`、`load`、`save`、warmup、optimizer / scheduler 参数等。
-
-如果你在 actor 和 critic 之间写入不同的并行拓扑，当前行为不受支持，可能导致初始化或训练过程出错。
-
----
-
-## FAQ
-
-### Q: 可以只写 actor 或只写 critic 吗？
-
-可以。缺失角色会自动继承公共 CLI 参数，不需要把所有参数都重复写一遍。
-
-### Q: 可以把 `--actor-num-nodes` 或 `--critic-num-gpus-per-node` 写进 YAML 吗？
-
-不可以。当前资源分配和 placement group 仍由 CLI 参数控制，YAML 中对应字段会被忽略。
\ No newline at end of file
diff --git a/docs/zh/advanced/sglang-config.md b/docs/zh/advanced/sglang-config.md
deleted file mode 100644
index e68c52c843..0000000000
--- a/docs/zh/advanced/sglang-config.md
+++ /dev/null
@@ -1,453 +0,0 @@
-# SGLang Config：高级引擎部署
-
-`--sglang-config` 是一个基于 YAML 的配置系统，用于在 slime 中精细控制 SGLang 引擎的部署。它支持**多模型服务**、**Prefill-Decode (PD) 分离**、**异构服务器组**，甚至可以作为复杂推理拓扑的**独立 SGLang 启动器**。
-
----
-
-## 架构概览
-
-在默认配置（不使用 `--sglang-config`）下，slime 部署单个模型，放在单个 router 后面，使用统一的服务器组：
-
-![架构概览](../../_static/image/arch.png)
-
-使用 `--sglang-config` 后，SGLang 部署扩展为多模型、多 router 拓扑：
-
-![sglang-config 架构](../../_static/image/sglang_config.png)
-
-**核心设计原则：**
-
-- **每个模型拥有独立的 router。** 模型在路由层隔离，支持独立的负载均衡和容错。
-- **同一模型内的服务器组可以异构。** 不同组可以有不同的 TP 大小、worker 类型（prefill/decode/regular）和 SGLang server 参数覆盖。
-- **权重同步按模型维度。** 只有 `update_weights: true` 的模型会接收来自训练的权重更新。冻结的模型（reference、reward 等）保持原样。
-
----
-
-## 配置格式
-
-配置文件是一个 YAML 文档，顶层 `sglang` 键包含一个模型定义列表：
-
-```yaml
-sglang:
-  - name: <model_name>              # 必填。模型的唯一标识符。
-    model_path: <path>              # 可选。HF checkpoint 路径。默认使用 --hf-checkpoint。
-    update_weights: <bool>          # 可选。是否从训练同步权重。自动推断。
-    num_gpus_per_engine: <int>      # 可选。该模型所有组的默认 TP 大小。
-    server_groups:                  # 必填。服务器组配置列表。
-      - worker_type: <type>         # 必填。可选：regular、prefill、decode、placeholder。
-        num_gpus: <int>             # 必填。分配给该组的 GPU 总数。
-        num_gpus_per_engine: <int>  # 可选。该组的 TP 大小覆盖。
-        overrides: <dict>           # 可选。SGLang ServerArgs 字段覆盖。
-```
-
-### 字段参考
-
-#### 模型级字段
-
-| 字段 | 类型 | 默认值 | 说明 |
-|------|------|--------|------|
-| `name` | `str` | **必填** | 模型唯一名称（如 `"actor"`、`"ref"`、`"reward"`）。用作 `args.sglang_model_routers` 的 key。 |
-| `model_path` | `str` | `args.hf_checkpoint` | HuggingFace checkpoint 路径。同一模型内的所有服务器组必须使用相同的 model path。 |
-| `update_weights` | `bool` | 自动推断 | 该模型是否接收训练权重更新。未设置时自动推断：如果 `model_path` 与 `--hf-checkpoint` 匹配则为 `true`，否则为 `false`。 |
-| `num_gpus_per_engine` | `int` | `args.rollout_num_gpus_per_engine` | 该模型服务器组的默认 TP 大小。各组可单独覆盖。 |
-| `server_groups` | `list` | **必填** | `ServerGroupConfig` 条目列表，定义引擎拓扑。（`engine_groups` 作为向后兼容别名仍可使用。） |
-
-#### 服务器组级字段
-
-| 字段 | 类型 | 默认值 | 说明 |
-|------|------|--------|------|
-| `worker_type` | `str` | **必填** | 引擎类型：`regular`（标准）、`prefill`（PD prefill worker）、`decode`（PD decode worker）或 `placeholder`（占位，不启动引擎）。 |
-| `num_gpus` | `int` | **必填** | 该组的 GPU 总数。必须 > 0。 |
-| `num_gpus_per_engine` | `int` | 模型的 `num_gpus_per_engine` | TP 大小覆盖。每个引擎实例的 GPU 数量。 |
-| `overrides` | `dict` | `{}` | SGLang `ServerArgs` 字段覆盖。优先级最高，覆盖 `--sglang-*` CLI 参数和模型级默认值。 |
-
-### Worker 类型
-
-| 类型 | 说明 | 使用场景 |
-|------|------|----------|
-| `regular` | 标准 SGLang 引擎 | 默认模式，同时处理 prefill 和 decode |
-| `prefill` | PD 分离的 prefill worker | 专门处理 prompt；与 `decode` worker 配对 |
-| `decode` | PD 分离的 decode worker | 专门生成 token；与 `prefill` worker 配对 |
-| `placeholder` | 占位，不创建引擎 | 为训练共置预留 GPU 或留作未来使用 |
-
----
-
-## 使用模式
-
-### 1. 基本单模型部署
-
-最简单的配置，复现默认行为：
-
-```yaml
-# sglang_basic.yaml
-sglang:
-  - name: default
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_basic.yaml \
-  --rollout-num-gpus 8 \
-  --rollout-num-gpus-per-engine 2 \
-  ...
-```
-
-这将创建 4 个引擎（8 GPU ÷ 2 GPU/引擎），位于单个 router 之后。
-
-### 2. PD 分离
-
-将 prefill 和 decode 阶段分离到专用服务器组，以提升多轮和 agentic 场景的吞吐量：
-
-```yaml
-# sglang_pd.yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2    # 2 个 prefill 引擎，TP=2
-      - worker_type: decode
-        num_gpus: 12
-        num_gpus_per_engine: 4    # 3 个 decode 引擎，TP=4
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_pd.yaml \
-  --rollout-num-gpus 16 \
-  ...
-```
-
-**为什么需要 PD 分离？** 在多轮场景中，prefill 和 decode 具有不同的计算特性。Prefill 是计算密集型（处理整个 prompt），而 decode 是内存带宽密集型（逐 token 生成）。分离它们可以：
-- 为 prefill 使用更小的 TP（提高每 GPU 吞吐量）
-- 为 decode 使用更大的 TP（降低延迟）
-- 独立扩展 prefill 和 decode 的容量
-
-> **注意：** PD 分离使用 SGLang Model Gateway (sgl-router)，并设置 `pd_disaggregation=True`。
-
-### 3. 多模型服务
-
-同时部署多个模型，每个模型拥有独立的 router：
-
-```yaml
-# sglang_multi_model.yaml
-sglang:
-  - name: actor
-    update_weights: true              # 接收训练权重更新
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: ref
-    model_path: /path/to/ref_model    # 不同的模型 checkpoint
-    update_weights: false              # 冻结，不更新权重
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-
-  - name: reward
-    model_path: /path/to/reward_model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-```bash
-python train.py \
-  --sglang-config sglang_multi_model.yaml \
-  --rollout-num-gpus 16 \
-  --hf-checkpoint /path/to/actor_model \
-  --rollout-function-path my_rollout.generate_rollout \
-  ...
-```
-
-**在自定义 rollout 函数中访问模型：**
-
-```python
-from slime.rollout.sglang_rollout import get_model_url
-from slime.utils.http_utils import post
-
-async def my_generate(args, sample, sampling_params):
-    # 路由到 actor 模型（默认）
-    actor_url = get_model_url(args, "actor", "/generate")
-    output = await post(actor_url, {"text": sample.prompt, "sampling_params": sampling_params})
-    
-    # 路由到 reference 模型
-    ref_url = get_model_url(args, "ref", "/generate")
-    ref_output = await post(ref_url, {"text": sample.prompt, "sampling_params": sampling_params})
-    
-    # 路由到 reward 模型（如 OpenAI 兼容 API）
-    reward_url = get_model_url(args, "reward", "/v1/chat/completions")
-    reward_output = await post(reward_url, {...})
-    
-    ...
-```
-
-`get_model_url()` 从 `args.sglang_model_routers`（一个将模型名称映射到 `(ip, port)` 元组的字典）中读取，该字典在引擎启动后自动填充。
-
-### 4. 多模型 + PD 分离
-
-将多模型与 PD 分离结合，实现最大灵活性：
-
-```yaml
-# sglang_full.yaml
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2
-      - worker_type: decode
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: ref
-    model_path: /path/to/ref_model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-### 5. 占位组用于 GPU 预留
-
-使用 `placeholder` 组来预留 GPU 而不创建引擎。这在共置训练场景中很有用，部分 GPU 需要为训练预留：
-
-```yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: regular
-        num_gpus: 6
-        num_gpus_per_engine: 2
-      - worker_type: placeholder
-        num_gpus: 2                   # 预留 2 个 GPU（不创建引擎）
-```
-
-### 6. 按组覆盖 ServerArgs
-
-使用 `overrides` 将 SGLang `ServerArgs` 字段应用到特定服务器组，而不影响其他组：
-
-```yaml
-sglang:
-  - name: actor
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-        overrides:
-          mem_fraction_static: 0.85
-          context_length: 32768
-          chunked_prefill_size: 4096
-          enable_torch_compile: true
-```
-
-覆盖具有**最高优先级**，会覆盖基础的 `--sglang-*` CLI 参数和模型级默认值。这对以下场景特别有用：
-- 不同组使用不同的内存配置
-- prefill 和 decode 使用不同的 context length
-- 在特定组上启用实验性功能
-
-### 7. 独立 SGLang 启动器
-
-虽然 `--sglang-config` 是为 slime 的训练流水线设计的，但它也可以作为纯推理场景的强大启动器，通过 `--rollout-external` 模式或配置 slime 仅关注推理服务。
-
-**使用预启动的外部引擎：**
-
-对于复杂的生产部署，你可能希望独立预启动 SGLang 引擎，然后将其连接到 slime：
-
-```bash
-# 步骤 1：外部启动 SGLang 引擎
-python -m sglang.launch_server --model-path /path/to/model --port 10090 ...
-python -m sglang.launch_server --model-path /path/to/model --port 10091 ...
-
-# 步骤 2：将 slime 连接到外部引擎
-python train.py \
-  --rollout-external \
-  --rollout-external-engine-addrs host1:10090 host2:10091 \
-  ...
-```
-
-> **注意：** `--sglang-config` 和 `--rollout-external` 互斥。当你希望 slime 管理完整的引擎生命周期时，使用 `--sglang-config`；当引擎已预部署时，使用 `--rollout-external`。
-
----
-
-## Router 配置
-
-每个模型都有独立的 router（默认使用 SGLang Model Gateway）。
-
-### Router 策略
-
-你可以配置路由策略：
-
-```bash
---router-policy round_robin        # 简单轮询
---router-policy consistent_hashing  # 多轮会话亲和
---router-policy cache_aware         # 缓存感知路由（默认）
-```
-
-### 多轮 Agent 的会话亲和路由
-
-对于多轮对话和 agentic 场景，会话亲和确保同一对话的所有请求路由到同一个 backend worker。这可以显著提升 prefix cache 命中率，因为 worker 已经缓存了对话历史。
-
-slime 自动为每个 sample 分配一个唯一的 `session_id`（存储在 `sample.session_id` 中）。当 router 策略为 `consistent_hashing` 时，该 ID 通过 `X-SMG-Routing-Key` header 传递，SGLang Model Gateway 使用它将同一会话的所有轮次确定性地路由到同一个 worker。
-
-```bash
---router-policy consistent_hashing
-```
-
-**工作原理：**
-
-1. 每个 sample 通过 UUID 分配唯一的 `session_id`
-2. 每次请求时，slime 在 HTTP header 中传递 `X-SMG-Routing-Key: <session_id>`
-3. SGLang Model Gateway 的 consistent hashing 策略将该 key 映射到特定的 worker
-4. 后续轮次复用相同的 `session_id`，确保命中同一个 worker
-
----
-
-## 解析规则
-
-加载配置时，slime 按以下优先级顺序解析：
-
-1. **每引擎 GPU 数回退：** 组 `num_gpus_per_engine` → 模型 `num_gpus_per_engine` → `args.rollout_num_gpus_per_engine`
-2. **模型路径回退：** 组 `overrides.model_path` → 模型 `model_path` → `args.hf_checkpoint`
-3. **权重更新推断：** 如果 `update_weights` 未设置：
-   - 如果有效的 model path 与 `--hf-checkpoint` 匹配，则为 `true`
-   - 否则为 `false`（会有警告）
-4. **GPU 总数校验：** 所有模型所有组的 `num_gpus` 总和必须等于 `--rollout-num-gpus`。
-
----
-
-## 互斥
-
-`--sglang-config` 与以下选项互斥：
-
-| 选项 | 冲突原因 |
-|------|----------|
-| `--prefill-num-servers` | PD 分离通过 YAML 中的 `server_groups` 配置 |
-| `--rollout-external` | 外部引擎有自己的拓扑；config 在内部管理生命周期 |
-
----
-
-## 完整示例：多模型 Agentic 训练
-
-下面是一个完整的示例，展示在 32 个 GPU 上使用 PD 分离进行多模型 agentic RL 训练：
-
-**配置文件 (`sglang_agent.yaml`)：**
-
-```yaml
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: prefill
-        num_gpus: 4
-        num_gpus_per_engine: 2
-        overrides:
-          chunked_prefill_size: 8192
-      - worker_type: decode
-        num_gpus: 12
-        num_gpus_per_engine: 4
-        overrides:
-          mem_fraction_static: 0.88
-
-  - name: ref
-    model_path: /data/models/Qwen3-32B
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-
-  - name: reward
-    model_path: /data/models/reward-model
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-```
-
-**启动命令：**
-
-```bash
-python train.py \
-  --sglang-config sglang_agent.yaml \
-  --hf-checkpoint /data/models/Qwen3-8B \
-  --rollout-num-gpus 32 \
-  --rollout-function-path my_agent.rollout.generate_rollout \
-  --custom-rm-path my_agent.reward.reward_func \
-  --advantage-estimator grpo \
-  --n-samples-per-prompt 8 \
-  ...
-```
-
-**自定义 rollout 函数 (`my_agent/rollout.py`)：**
-
-```python
-from slime.rollout.sglang_rollout import get_model_url
-from slime.utils.http_utils import post
-
-async def generate_with_models(args, sample, sampling_params):
-    """使用 actor 生成，用 reward 模型打分，与 reference 比较。"""
-    
-    # 从 actor 生成
-    actor_url = get_model_url(args, "actor", "/generate")
-    actor_output = await post(actor_url, {
-        "text": sample.prompt,
-        "sampling_params": sampling_params,
-        "return_logprob": True,
-    })
-    
-    # 获取 reference logprobs 用于 KL penalty
-    ref_url = get_model_url(args, "ref", "/generate")
-    ref_output = await post(ref_url, {
-        "text": sample.prompt + actor_output["text"],
-        "sampling_params": {"max_new_tokens": 0, "temperature": 0},
-        "return_logprob": True,
-    })
-    
-    # 用 reward 模型打分
-    reward_url = get_model_url(args, "reward", "/v1/chat/completions")
-    reward_output = await post(reward_url, {
-        "model": "reward",
-        "messages": [{"role": "user", "content": sample.prompt + actor_output["text"]}],
-    })
-    
-    # ... 处理输出并返回 Sample
-```
-
----
-
-## FAQ
-
-### Q: 同一模型内可以混用 PD 和 regular 组吗？
-
-不可以。PD 分离要求一个模型的服务器组要么全部是 prefill/decode 对，要么全部是 regular。不支持在同一模型内混用 `regular` 与 `prefill`/`decode`。
-
-### Q: 如果 `num_gpus` 不能被 `num_gpus_per_engine` 整除怎么办？
-
-对于跨节点引擎（`num_gpus_per_engine > num_gpus_per_node`），划分基于每节点的本地 GPU 数量。例如，每节点 8 个 GPU 且 `num_gpus_per_engine: 16` 时，每个引擎横跨 2 个节点。
-
-### Q: 同一模型内的不同服务器组可以使用不同的 model path 吗？
-
-不可以。同一模型内的所有服务器组必须共享相同的 `model_path`。这在 `resolve()` 时会被校验。如果需要不同的模型，请定义为独立的模型条目。
-
-### Q: 运行时如何获取特定模型的 router 地址？
-
-使用 `slime.rollout.sglang_rollout` 中的 `get_model_url(args, "model_name", "/endpoint")`。它从 `args.sglang_model_routers`（一个 `{ model_name: (ip, port) }` 字典）中读取，该字典在引擎启动后自动填充。
-
-### Q: 可以不训练，只用 `--sglang-config` 做推理吗？
-
-虽然 `--sglang-config` 是为 slime 的训练循环设计的，但你可以通过配置仅 rollout 的运行来实现纯推理场景。对于完全独立的 SGLang 推理服务，建议直接使用 SGLang 原生的 `launch_server`，或使用 `--rollout-external` 模式连接预部署的引擎。
-
-### Q: `--sglang-config` 和 `--prefill-num-servers` 是什么关系？
-
-`--prefill-num-servers` 是启用 PD 分离的旧方式（它创建一个带有 prefill + decode 组的单模型）。`--sglang-config` 是更新、更灵活的方式。两者互斥。我们推荐所有新部署迁移到 `--sglang-config`。
diff --git a/docs/zh/advanced/slime-router.md b/docs/zh/advanced/slime-router.md
new file mode 100644
index 0000000000..c7da222e8a
--- /dev/null
+++ b/docs/zh/advanced/slime-router.md
@@ -0,0 +1,138 @@
+# Slime Router
+
+Slime 提供一个可选的 SlimeRouter，用于 rollout / data generation 阶段。它是一个轻量级的 HTTP router/proxy，位于一个或多个 SGLang worker server 前，补齐一些 training-oriented 能力——这些并不是 serving-focused router 的主要目标。
+
+---
+
+## 1. 什么是 SlimeRouter？
+
+SlimeRouter 是一个小型 FastAPI 服务，主要能力包括：
+
+- 注册 worker（SGLang HTTP server）到本地池
+- 路由请求到选定的 worker（简单的 least-inflight load balancing）
+- 代理任意路径到选定的 worker（例如 `/generate`）
+- 定期 health checks，并隔离不健康的 worker
+- 支持 middleware plugins（通过 `--slime-router-middleware-paths`）实现 rollout 特定处理（例如 caching、request/response transform）
+
+在 slime 架构中，router 是 rollout 系统（"SGLang + router"）的一部分：负责生成样本并将其推入数据缓冲区。
+
+### 启动方式
+
+在分布式训练中，当未提供 `--sglang-router-ip` 时，slime 会自动启动一个 router：
+
+- 如果设置了 `--use-slime-router`，slime 启动 SlimeRouter
+- 否则，slime 启动 SGLang Model Gateway
+
+---
+
+## 2. 为什么需要 SlimeRouter
+
+与 production inference 不同，RL rollout 往往需要捕获用于训练的额外 metadata：token-level logprobs、loss masks，以及（对 MoE 模型）expert routing decisions。SlimeRouter 通过 middleware system 和 passthrough proxy 设计提供这些能力。
+
+### 2.1 Radix-tree cache（透明的 token 管理）
+
+> 当你的 rollout 流程是 text-in/text-out、并且很难可靠地保存 token IDs 时，适合用 radix-tree cache；如果你已经能自己控制 token-in/token-out（例如 search r1、multiturn VLM 这些 example），通常不需要 radix-tree cache。
+
+text-in text-out 接口可能导致 token retokenization mismatches：训练阶段重新 tokenize 文本，得到的 token 序列可能与 rollout 阶段不同，从而破坏 PPO/GRPO 这类方法所需的 per-token alignment。
+
+radix-tree cache 可以透明地解决这个问题：它拦截 text-based request，对其进行 tokenize，并将 trajectory（text、token IDs、logprobs、loss masks）按文本前缀作为 key 存储。rollout 结束后，调用 `/retrieve_from_text` 就能取回与 rollout 完全一致的 token 序列及其对齐的 metadata，无需修改现有 rollout 代码。
+
+实现上，radix-tree cache 会做几件事：
+
+- 接收 text 以及 tokens/metadata，并写入 radix tree
+- 通过 longest-prefix matching 复用已缓存的 token 序列（使后续流程可以走 token-in/token-out）
+- rollout 过程中持续插入新的 text continuation（同一 prompt 下可有多条 trajectory，例如 GRPO）
+- 定期清理 stale nodes，控制内存占用
+
+当你有 text-based rollout 代码、想获得 token-level 精度但又不想重写，或者在 GRPO 场景中多个 trajectory 共享相同 prompt 前缀时，建议使用 radix-tree cache。
+
+### 2.2 Rollout routing replay (R3) for MoE
+
+对 MoE 模型，slime 支持 rollout routing replay (R3)：在 rollout 期间记录 expert routing decisions，并在训练期间 replay，以提升训练稳定性。
+
+#### SGLang 端
+
+SGLang 侧通过以下机制提供 expert routing capture：
+
+- `--enable-return-routed-experts`：启用路由捕获的服务器参数
+- `RoutedExpertsCapturer`：在前向传播期间捕获每个 MoE 层的 `topk_ids`（选定的专家 ID）
+- `return_routed_experts`：检索路由数据的请求参数
+- 在响应 `meta_info` 中返回 `routed_experts`：一个形状为 `[seq_len - 1, num_layers, top_k]` 的 expert ID tensor
+
+#### Slime 端
+
+Slime 侧消费路由数据，并在训练中完成 replay：
+
+- `--use-slime-router --use-rollout-routing-replay`：启用 R3 需要同时设置这两个标志
+- Rollout 发送 `return_routed_experts=True` 并将结果存储在 `sample.rollout_routed_experts` 中
+- 训练调用 `fill_routing_replay()` 将路由数据加载到 `RoutingReplay` 对象中
+- forward pass 期间，直接 replay 记录的 routing decisions，而不是重新计算
+
+#### 为什么需要 SlimeRouter
+
+我们需要 SlimeRouter，是因为当请求设置 `return_routed_experts=true` 时，SGLang worker 会在响应里返回路由信息（`meta_info.routed_experts`），而 SlimeRouter 会端到端保留这个字段。SGLang Model Gateway 会用固定 schema 重建响应，可能会丢掉这类额外 metadata（细节见第 3 节）。
+
+---
+
+## 3. 与 SGLang Model Gateway 的区别
+
+SlimeRouter 与 SGLang Model Gateway 都能将请求路由到 worker，但它们面向的目标不同、优化方向也不同。
+
+### 主要区别
+
+SlimeRouter 是一个轻量级的 Python/FastAPI proxy，作为 SGLang worker 的 passthrough proxy。这种 passthrough 设计使得 RL 特定功能成为可能，例如 radix-tree trajectory caching 和 R3（需要保留原始 response metadata，如 `routed_experts`）。
+
+SGLang Model Gateway 是一个高性能 Rust router，面向大规模 inference 优化：async non-blocking routing、高级 fault tolerance（retries、circuit breakers）、多种 load balancing policy（包括 cache-aware routing），以及 PD disaggregation 支持。但它会用固定 schema 重建响应，因此不保留 slime 的 R3 流程所需 metadata。
+
+更多关于 SGLang Model Gateway 的信息，请参阅[官方文档](https://docs.sglang.io/advanced_features/sgl_model_gateway.html)。
+
+### 如何选择
+
+- 当你需要 R3 或 radix-tree cache 时，使用 SlimeRouter
+- 其他情况使用 SGLang Model Gateway（推荐默认选项）
+
+---
+
+## 4. 多轮 Agent 的会话亲和性路由
+
+当使用 SGLang Model Gateway 的一致性哈希路由策略时，slime 会自动为每个 rollout session 分配唯一的 session ID，并将其作为路由键来实现会话亲和性。
+
+### 什么是会话亲和性？
+
+会话亲和性（Session Affinity，也称为粘性会话）确保属于同一对话或 agent session 的所有请求都被路由到相同的后端 worker。这对以下场景有益：
+
+- **多轮对话**：保持相同 worker 可以提高前缀缓存命中率
+- **多 agent 系统**：确保 agent 状态一致性和更好的资源局部性
+- **调试**：更容易追踪和调试特定会话
+
+### 工作原理
+
+当 rollout 系统生成样本时，每个样本都会被分配一个唯一的 `session_id`。这个 ID 会：
+
+1. 使用 UUID 为每个样本自动生成
+2. 存储在 `sample.session_id` 字段中
+3. 当 router policy 为 `consistent_hashing` 时，作为 `X-SMG-Routing-Key` header 传递
+
+SGLang Model Gateway 的一致性哈希策略会使用这个路由键，确定性地为所有具有相同 session ID 的请求选择相同的 worker。
+
+### 配置方法
+
+启用会话亲和性路由只需在 slime 中配置 router policy 参数：
+
+```bash
+--sglang-router-policy consistent_hashing
+```
+
+Slime 会自动启动 SGLang Model Gateway 并使用一致性哈希策略。
+
+> **注意**：如果遇到 `consistent_hashing` policy 不可用的错误，请升级 sglang-router：
+> ```bash
+> pip install -U sglang-router
+> ```
+
+### 注意事项
+
+- 每个样本都有自己独立的 session ID
+- 同一 group 中的不同样本可能会被路由到不同的 worker
+- 同一样本的后续轮次会保持相同的 session ID
+- 目前该功能只适用于 SGLang Model Gateway
diff --git a/docs/zh/developer_guide/ci.md b/docs/zh/developer_guide/ci.md
index b8830ecc77..a43590d939 100644
--- a/docs/zh/developer_guide/ci.md
+++ b/docs/zh/developer_guide/ci.md
@@ -6,7 +6,7 @@ slime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发——
 
 工作流定义在 `.github/workflows/pr-test.yml`（由 `pr-test.yml.j2` 自动生成）。每个 CI 任务会：
 
-1. 在自托管 GPU runner 上通过 `docker run` 运行；大多数测试使用 `slimerl/slime:latest`，镜像验证使用 `slimerl/slime-test:latest`。
+1. 在自托管 GPU runner 上以 Docker 容器（`slimerl/slime:latest`）运行。
 2. 通过 `pip install -e . --no-deps` 安装 slime。
 3. 通过 `tests/ci/gpu_lock_exec.py --count <num_gpus>` 获取所需数量的 GPU。
 4. 执行测试文件：`python <test_path>.py` 或 `python tests/<test_file>.py`。如果测试位于 `tests/plugin_contracts/` 这样的子目录，CI 也会自动处理。
@@ -20,8 +20,9 @@ slime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发——
 | Label | Job | 说明 |
 |---|---|---|
 | `run-ci-short` | `e2e-test-short` | Qwen2.5-0.5B 轻量级冒烟测试（4 GPU），用于快速反馈。 |
+| `run-ci-fsdp` | `e2e-test-fsdp` | FSDP 后端测试（true on-policy、VL、megatron-fsdp 对齐）。 |
 | `run-ci-megatron` | `e2e-test-megatron` | 核心 Megatron 训练测试，覆盖 Dense、MoE、PPO、MTP、OPD 等。 |
-| `run-ci-precision` | `e2e-test-precision` | 数值精度校验（并行一致性检查）。 |
+| `run-ci-precision` | `e2e-test-precision` | 数值精度校验（并行一致性检查、megatron-fsdp 对齐）。 |
 | `run-ci-ckpt` | `e2e-test-ckpt` | Checkpoint 保存/加载正确性（同步和异步保存）。 |
 | `run-ci-image` | `e2e-test-image` | 在 `slimerl/slime-test:latest` 镜像上运行**全部**测试（用于镜像验证）。 |
 | `run-ci-changed` | `e2e-test-changed` | **动态**检测 PR 中新增或修改的测试文件，仅运行这些测试。 |
@@ -56,7 +57,7 @@ slime 使用 GitHub Actions 进行 CI。测试通过 **PR label** 触发——
 这是验证 Megatron 后端改动的主要 label，覆盖：
 
 - Dense 模型：GLM4-9B、Qwen3-4B（PPO）
-- MoE 模型：Qwen3-30B-A3B（DeepEP + FP8）、Qwen3.6-35B-A3B PD + Mooncake、Moonlight-16B-A3B
+- MoE 模型：Qwen3-30B-A3B（有/无 DeepEP + FP8）、Moonlight-16B-A3B
 - 特殊场景：MiMo-7B MTP、Qwen2.5-0.5B debug rollout-then-train、OPD（sglang teacher 模式）
 
 所有测试使用 8 张 GPU。如果你正在修改 Megatron 训练逻辑、loss 计算或 checkpoint 转换，应该使用这个 label。
@@ -75,7 +76,7 @@ NUM_GPUS = 4  # 此常量会被 run-ci-changed 自动读取
 
 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.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     # 按需下载数据集 ...
 
 def execute():
diff --git a/docs/zh/developer_guide/trace.md b/docs/zh/developer_guide/trace.md
deleted file mode 100644
index 09588bd32a..0000000000
--- a/docs/zh/developer_guide/trace.md
+++ /dev/null
@@ -1,119 +0,0 @@
-# Trace 可视化
-
-slime 可以为每条 rollout sample 挂上轻量级执行 trace。它会记录生成、奖励模型等 span 事件，并且可以在保存下来的 rollout debug dump 中离线查看。
-
-![trace 时间线查看器](../../_static/image/trace.png)
-
-## 保存 rollout trace 数据
-
-如果想在运行结束后查看 trace，可以在训练时打开 rollout debug dump：
-
-```bash
-python train.py \
-    ... \
-    --save-debug-rollout-data /path/to/debug/rollout_{rollout_id}.pt
-```
-
-每个保存出来的 `.pt` 文件都会包含 rollout samples，以及对应的 `trace` 数据。之后也可以通过 `--load-debug-rollout-data` 复用同一份 dump。
-
-## 打开时间线查看器
-
-对保存好的 rollout dump 运行：
-
-```bash
-python tools/trace_timeline_viewer.py /path/to/debug/rollout_0.pt
-```
-
-脚本会生成：
-
-- `rollout_0.trace_timeline_cache.json`
-- `rollout_0.trace_timeline_viewer.html`
-
-默认情况下，它还会启动一个本地静态文件服务，方便直接在浏览器里打开。如果只想生成文件，可以加 `--no-serve`。
-
-## 如何理解可视化结果
-
-- 每一行对应一条 sample。
-- 条形块表示 span，点表示瞬时事件。
-- `trace_span(...)` 在开始和结束时记录的属性，都会显示在详情面板里。
-- 当 SGLang 返回 PD 分离相关时延时，viewer 会自动补出 `[P]` 和 `[D]` 两条虚拟 lane，用来拆开展示 prefill/decode。
-- 如果没有开启 PD，这两条虚拟 lane 不会出现，基础 trace 也仍然可以正常渲染。
-
-## 给自定义代码打点
-
-在自定义 rollout 或 reward 逻辑中，可以直接复用 `slime.utils.trace_utils` 里的工具：
-
-- `trace_span(target, name, attrs=...)`：记录一段持续时间。
-- `trace_event(target, name, attrs=...)`：记录一个瞬时事件。
-- `trace_function(name, ...)`：用 decorator 把整个同步/异步函数包成一个 span。
-- `bind_trace(sample)`：在 sample 被传递到其他 helper 或任务之前，确保它已经绑定好 trace carrier。
-
-### `trace_span` 和 `trace_function` 的区别
-
-如果你只想给函数内部的一小段逻辑打点，或者需要在代码块执行过程中补充 span 结束属性，就用 `trace_span(...)`。
-
-如果整个函数调用都应该对应一个 span，就用 `trace_function(...)`。它本质上会先解析 trace target，然后在函数调用外层自动套一层 `trace_span(...)`，因此同步函数和异步函数都能直接用。
-
-slime 主 rollout 流程里就是这样用的。例如 `generate_and_rm(...)` 按 sample 打点，而 `generate_and_rm_group(...)` 按 group 打点：
-
-```python
-from slime.utils.trace_utils import trace_function
-
-
-@trace_function("generate_and_rm", target="sample")
-async def generate_and_rm(args, sample, sampling_params, evaluation=False):
-    ...
-
-
-@trace_function(
-    "generate_and_rm_group",
-    target="group",
-    attrs_getter=lambda args, group, sampling_params, evaluation=False: {"group_size": len(group)},
-)
-async def generate_and_rm_group(args, group, sampling_params, evaluation=False):
-    ...
-```
-
-### 如何选择 target
-
-`trace_function(...)` 需要一个 trace target，通常是 `Sample`、`TraceHandle` 或它们的列表。
-
-- 如果 target 本来就是函数参数，优先写成 `target="sample"` 或 `target="group"`。
-- 如果 target 需要从参数里推导出来，就用 `target_getter=...`。
-- 不建议默认依赖自动推断。实现里确实可以从参数或当前 trace 上下文里推断 target，但显式写出来更稳定，也更不容易出现歧义。
-
-### 给 decorator span 增加属性
-
-如果想在 span 开始时附带属性，可以用 `attrs_getter=...`：
-
-```python
-@trace_function(
-    "custom_rollout_batch",
-    target="samples",
-    attrs_getter=lambda samples, **_: {"batch_size": len(samples)},
-)
-async def custom_rollout_batch(samples, **kwargs):
-    ...
-```
-
-如果你需要在函数执行到一半之后再补充属性，那么不要只靠 decorator，应该在函数内部再嵌一层 `trace_span(...)`。一个比较常见的模式是：
-
-- 外层用 `trace_function(...)` 表示整个函数生命周期
-- 内层用 `trace_span(...)` 标记 generation、RM、filter、post-process 等关键子步骤
-
-如果想统一记录 SGLang 返回的 generation 元信息，可以复用 `build_sglang_meta_trace_attrs`：
-
-```python
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
-
-with trace_span(sample, "sglang_generate") as span:
-    output = await post(url, payload)
-    span.update(build_sglang_meta_trace_attrs(output["meta_info"]))
-```
-
-## 使用建议
-
-- 先保存少量 rollout；单个 dump 的 sample 数量适中时，viewer 会更容易阅读。
-- viewer 直接基于保存下来的 `.pt` dump 工作，因此可以把文件拷到别的机器离线分析。
-- 如果你想看的是 SGLang 自身的 GPU / kernel 级 profiling trace，请参考 [性能分析](./profiling.md)。
-
diff --git a/docs/zh/examples/deepseek-r1.md b/docs/zh/examples/deepseek-r1.md
index 703a7730e4..368653844d 100644
--- a/docs/zh/examples/deepseek-r1.md
+++ b/docs/zh/examples/deepseek-r1.md
@@ -171,7 +171,7 @@ OPTIMIZER_ARGS=(
 
 sglang 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 sglang 的 `tp_size`，除此之外的 sglang 参数均通过添加 `--sglang-` 的前缀来传给 slime。为了充分利用 sglang 的大 EP 推理能力，我们加上了 ep64、dp_attention dp8、deepep mode auto 等配置。
 
-最后的 `--rollout-server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。
+最后的 `--sglang-server-concurrency` 是 slime 的特有参数，是为了方式同时发给 sglang server 的并发太大打爆 http server，默认为 512。但是我们现在是 8 机一个 server，为了保证每个 dp rank 能有 128 的并发，我们调整为 1024。
 
 ```bash
 SGLANG_ARGS=(
@@ -190,7 +190,7 @@ SGLANG_ARGS=(
    --sglang-deepep-mode auto
 
    # make every dp rank has 128 concurrency
-   --rollout-server-concurrency 1024
+   --sglang-server-concurrency 1024
 )
 ```
 
diff --git a/docs/zh/examples/glm4.5-355B-A32B.md b/docs/zh/examples/glm4.5-355B-A32B.md
new file mode 100644
index 0000000000..8a9db6d698
--- /dev/null
+++ b/docs/zh/examples/glm4.5-355B-A32B.md
@@ -0,0 +1,190 @@
+# 64xH100 训练 GLM-4.5
+
+这里是使用 64xH100 进行 GLM-4.5 355B RL 训练的示例。
+
+## 环境准备
+
+搭建环境与下载数据的方法可以参考 [示例：Qwen3-4B](qwen3-4B.md)。
+
+首先需要在多机均可访问到的地址（下记为 `$BASE_DIR`）上下载 GLM-4.5：
+
+```bash
+hf download zai-org/GLM-4.5 --local-dir $BASE_DIR/GLM-4.5-355B-A32B
+```
+
+通过如下方式通过 2 机 16 卡将 huggingface checkpoint 转换为 torch dist 格式：
+
+```bash
+cd slime/
+source scripts/models/glm4.5-355B-A32B.sh
+PYTHONPATH=/root/Megatron-LM/ torchrun \
+   --nproc-per-node 8 \
+   --master-addr ${MASTER_ADDR} --master-port 12345 \
+   --nnodes=2 --node-rank ${NODE_RANK} \
+   tools/convert_hf_to_torch_dist.py \
+   ${MODEL_ARGS[@]} \
+   --hf-checkpoint $BASE_DIR/GLM-4.5-355B-A32B/ \
+   --save $BASE_DIR/GLM-4.5-355B-A32B_torch_dist/
+```
+
+其中 `MASTER_ADDR` 为 node0 的 ip，`NODE_RANK` 表示这是第几台机器，这两者就像是在多机 `torchrun` 的时候进行的配置。
+
+## 执行训练
+
+在 node0 运行：
+
+```bash
+cd slime/
+bash scripts/run-glm4.5-355B-A32B.sh
+```
+
+在其他 node 需要通过如下的指令加入 ray 集群：
+
+```bash
+ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats"
+```
+
+或者如果你能获取到所有节点的 ip 列表，例如有一个 mpi hostfie（每一行为 `ip slot=8`），那么可以在 `scripts/run-glm4.5-355B-A32B.sh` 中的 `ray start --head` 指令之后加入如下的指令，从而只需要从 node0 执行训练：
+
+```bash
+for WORKER_IP in $(awk '{print $1}' $BASE_DIR/mpi_hostfile); do
+  if [[ "$WORKER_IP" == "$MASTER_ADDR" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" &
+done
+wait
+```
+
+### 参数简介
+
+```bash
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/glm4.5-355B-A32B.sh"
+```
+
+从 [scripts/models/glm4.5-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/models/glm4.5-355B-A32B.sh) 读取模型的 config。这些 config 都是 megatron 的参数。在使用 megatron 进行训练的时候，megatron 无法从 ckpt 中读取模型 config，需要我们自行配置。我们在 [scripts/models](https://github.com/THUDM/slime/tree/main/scripts/models/) 中提供了一些样例。
+
+#### PERF_ARGS
+
+一堆 megatron 的并行参数，只有 `--use-dynamic-batch-size` 与 `--max-tokens-per-gpu` 是 slime 添加的。
+
+megatron 的部分，我们配置了 tp8、pp4、cp2、ep16。
+
+`max_tokens_per_gpu` 是指每张卡最多跑多少 token，在开启 `use_dynamic_batch_size` 之后，会尽可能将一个 batch 内部长短不一的数据拼到 `max_tokens_per_gpu`，从而组成动态的 micro batch size，如果有一条数据长度超过了 `max_tokens_per_gpu`，则自成一条，不会对数据进行截断。在开启 context parallel (CP) 时，会让 CP 张卡去上的数据去共享总长为 `CP * max_tokens_per_gpu` 的 token。
+
+在开启 dynamic_batch_size，会忽略传统的 `micro_batch_size`。
+
+⚠️  slime 总是会通过 data packing 的方法训练模型，并且严格保证 per sample loss 或 per token loss，也就是开启 dynamic batch size 不会对 loss 计算有影响，推荐开启。
+
+```bash
+PERF_ARGS=(
+   --tensor-model-parallel-size 8
+   --sequence-parallel
+   --pipeline-model-parallel-size 4
+   --context-parallel-size 2
+   --expert-model-parallel-size 16
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 16384
+)
+```
+
+#### GRPO_ARGS
+
+目前 slime 这是一些 grpo 相关的参数：
+
+```bash
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+```
+
+如果希望训练时不加载 reference model，需要去掉 `--use-kl-loss` 并设置 `--kl-coef 0.00`（默认值为 0）。
+
+#### OPTIMIZER_ARGS
+
+我们通过了如下几个参数配置了 CPU Adam，用来节省显存。
+
+```bash
+OPTIMIZER_ARGS=(
+   ...
+
+   --optimizer-cpu-offload
+   --overlap-cpu-optimizer-d2h-h2d
+   --use-precision-aware-optimizer
+)
+```
+
+#### SGLANG_ARGS
+
+sglang 所需的参数，这里 `--rollout-num-gpus-per-engine` 基本对应 sglang 的 `tp_size`，除此之外的 sglang 参数均通过添加 `--sglang-` 的前缀来传给 slime。
+
+```bash
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 32
+   --sglang-mem-fraction-static 0.7
+   --sglang-enable-dp-attention
+   --sglang-dp-size 4
+)
+```
+
+#### MISC_ARGS
+
+一些额外的 megatron 配置。注意这里配置了 megatron 的 deepep。
+
+```bash
+MISC_ARGS=(
+   ...
+
+   # use deepep for megatron
+   --moe-enable-deepep
+   --moe-token-dispatcher-type flex
+)
+```
+
+## 用 fp8 进行数据生成
+
+开源版本的 GLM-4.5 fp8 ckpt 使用的是 per-channel 量化，目前无法在 sglang 中使用 deepep。我们可以利用 slime 提供的工具转换一个 128x128 per-block 量化的 checkpoint：
+
+```bash
+cd slime/
+python tools/convert_hf_to_fp8.py \
+    --model-dir $BASE_DIR/GLM-4.5-355B-A32B/ \
+    --save-dir $BASE_DIR/GLM-4.5-355B-A32B-FP8/ \
+    --strategy block --block-size 128 128 \
+    --max-workers 4
+```
+
+之后将 `--hf-checkpoint` 设置为 `$BASE_DIR/GLM-4.5-355B-A32B-FP8/` 就可以在训练中使用 fp8 进行 rollout 了。
+
+一个样例的 fp8 `SGLANG_ARGS` 为：
+
+```bash
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 32
+   --sglang-mem-fraction-static 0.7
+   --sglang-enable-dp-attention
+   --sglang-dp-size 32
+   --sglang-ep-size 32
+   --sglang-moe-dense-tp-size 1
+   --sglang-enable-dp-lm-head
+   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 128)
+
+   --sglang-moe-a2a-backend deepep
+   --sglang-deepep-mode auto
+)
+```
diff --git a/docs/zh/examples/glm4.7-30B-A3B.md b/docs/zh/examples/glm4.7-30B-A3B.md
index cdb2b45629..0a7a141c27 100644
--- a/docs/zh/examples/glm4.7-30B-A3B.md
+++ b/docs/zh/examples/glm4.7-30B-A3B.md
@@ -89,9 +89,9 @@ SGLANG_ARGS=(
    ...
    # MTP 投机解码 (EAGLE)
    --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
+   --sglang-speculative-num-steps 2
    --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --sglang-speculative-num-draft-tokens 3
 )
 ```
 
@@ -101,7 +101,7 @@ SGLANG_ARGS=(
 
 #### MTP 训练
 
-slime 也支持将 MTP 层与主模型联合训练，适用于已实现 MTP 权重转换的模型（如 MiMo、GLM-4.7）。启用时，相关参数如下：
+slime 也支持将 MTP 层与主模型联合训练，适用于已实现 MTP 权重转换的模型（如 MiMo、GLM-4.5）。启用时，相关参数如下：
 
 ```bash
 # 在模型配置中添加 MTP 层数
@@ -118,9 +118,9 @@ SPEC_ARGS=(
 - `--enable-mtp-training`：启用 MTP 层的梯度计算。不设置此标志时，MTP 层会被加载但冻结。
 - `--mtp-loss-scaling-factor 0.2`：MTP loss 相对于主策略 loss 的权重，默认为 0.2。
 
-> **注意**：MTP 训练需要 MTP checkpoint bridge 正确转换 HuggingFace 和 Megatron 格式之间的权重。`GLM4MoELiteBridge`（位于 `slime_plugins/mbridge/glm4moe_lite.py`）扩展了 DeepSeek V3 bridge，实现了动态 MTP 层索引以支持 GLM-4.7-Flash 的 47 层架构。
+> ⚠️ **注意**：GLM-4.7-Flash 的 MTP 训练目前尚不支持，因为 deepseek_v3 的 checkpoint bridge 尚未实现 MTP 权重转换（上游 mbridge 中标注为 `# TODO: mtp`）。但推理时的投机解码仍然可用 — SGLang 会内部处理 MTP 层。
 >
-> 对于其他支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。
+> 对于完整支持 MTP 训练的模型（如 MiMo），可参考 `scripts/run-mimo-7B-rl-eagle.sh`。
 
 ### 多机支持
 
diff --git a/docs/zh/examples/glm4.7-355B-A32B.md b/docs/zh/examples/glm4.7-355B-A32B.md
deleted file mode 100644
index 0954721da4..0000000000
--- a/docs/zh/examples/glm4.7-355B-A32B.md
+++ /dev/null
@@ -1,189 +0,0 @@
-# 64xH100 训练 GLM-4.7
-
-## 环境准备
-
-搭建环境与下载数据的方法与 Qwen3-4B 模型相同，可以参考 [示例：Qwen3-4B](qwen3-4B.md)，将文中 Qwen3-4B 的部分替换为 GLM-4.7 即可。
-
-### 前置条件
-
-GLM-4.7 使用 slime 标准 Docker 环境即可。多机启动前，请确保所有机器都能访问同一个 `$BASE_DIR` 路径，并在启动 Ray worker 前先取消代理：
-
-```bash
-unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
-```
-
-### 下载模型
-
-```bash
-hf download zai-org/GLM-4.7 --local-dir $BASE_DIR/GLM-4.7-355B-A32B
-```
-
-### 转换 Checkpoint
-
-可以用如下方法把 Hugging Face checkpoint 转换为 torch_dist 格式（2 机 x 8 卡）：
-
-```bash
-cd /root/slime
-pip install -e . --no-deps
-source scripts/models/glm4.5-355B-A32B.sh
-PYTHONPATH=/root/Megatron-LM/ torchrun \
-   --nproc-per-node 8 \
-   --master-addr ${MASTER_ADDR} --master-port 12345 \
-   --nnodes=2 --node-rank ${NODE_RANK} \
-   tools/convert_hf_to_torch_dist.py \
-   ${MODEL_ARGS[@]} \
-   --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B/ \
-   --save $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/
-```
-
-其中 `MASTER_ADDR` 是 node0 的 IP，`NODE_RANK` 表示当前机器的编号，配置方式与普通多机 `torchrun` 一致。
-
-## 执行训练
-
-从 node0 执行训练脚本：
-
-```bash
-cd /root/slime
-export BASE_DIR=/shared/path  # 所有节点都能访问的共享路径
-bash scripts/run-glm4.7-355B-A32B.sh
-```
-
-### 参数简介
-
-这里我们简单介绍一下 [run-glm4.7-355B-A32B.sh](https://github.com/THUDM/slime/blob/main/scripts/run-glm4.7-355B-A32B.sh) 中的关键部分。
-
-#### MoE 配置
-
-GLM-4.7 是一个 MoE（混合专家）模型，包含 160 个路由专家（top-8 激活）和共享专家。模型共 92 层：3 层 dense + 89 层 MoE。
-
-1. 为了支持在 64xH100 环境中运行 GLM-4.7，我们开启 Megatron 的 CPU Adam 来节省显存：
-
-   ```bash
-   OPTIMIZER_ARGS=(
-      ...
-      --optimizer-cpu-offload
-      --overlap-cpu-optimizer-d2h-h2d
-      --use-precision-aware-optimizer
-   )
-   ```
-
-2. 在 Megatron 中开启 MoE 优化。当前 64xH100 示例使用 TP=8、PP=4、CP=2、EP=16：
-
-   ```bash
-   PERF_ARGS=(
-      --tensor-model-parallel-size 8
-      --sequence-parallel
-      --pipeline-model-parallel-size 4
-      --context-parallel-size 2
-      --expert-model-parallel-size 16
-      --expert-tensor-parallel-size 1
-      ...
-      --use-dynamic-batch-size
-      --max-tokens-per-gpu 16384
-   )
-   ```
-
-3. 在 SGLang 中开启带 DP attention 的 MoE 优化：
-
-   ```bash
-   SGLANG_ARGS=(
-      --rollout-num-gpus-per-engine 32
-      --sglang-mem-fraction-static 0.7
-      --sglang-enable-dp-attention
-      --sglang-dp-size 4
-      --sglang-ep-size 32
-      --sglang-enable-dp-lm-head
-      --sglang-moe-dense-tp-size 1
-      ...
-   )
-   ```
-
-#### MTP 投机解码（推理加速）
-
-GLM-4.7 包含 MTP（Multi-Token Prediction）层，可以在推理阶段用于投机解码，加速 rollout 生成。启用方法是在 `SGLANG_ARGS` 中加入：
-
-```bash
-SGLANG_ARGS=(
-   ...
-   # MTP 投机解码 (EAGLE)
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
-)
-```
-
-这样 SGLang 就会使用模型自带的 MTP 层作为 EAGLE 风格投机解码的 draft model。
-
-> ⚠️ **注意**：投机解码会额外占用 GPU 显存。如果遇到 OOM，可以尝试降低 `--sglang-mem-fraction-static` 或暂时关闭投机解码。
-
-#### MTP 训练
-
-slime 也支持在 GLM-4.7 上将 MTP 层与主模型联合训练。启用时，相关参数如下：
-
-```bash
-# 在模型配置中添加 MTP 层数
-MODEL_ARGS+=(--mtp-num-layers 1)
-
-# 启用 MTP 训练
-MTP_ARGS=(
-   --enable-mtp-training
-   --mtp-loss-scaling-factor 0.2
-)
-```
-
-- `--mtp-num-layers 1`：告知 Megatron 从 checkpoint 中加载 MTP 层。
-- `--enable-mtp-training`：启用 MTP 层的梯度计算；不设置时 MTP 层会被加载但保持冻结。
-- `--mtp-loss-scaling-factor 0.2`：MTP loss 相对主策略 loss 的权重，默认值为 0.2。
-
-> **注意**：GLM-4.7 的 MTP 训练依赖 `GLM4MoEBridge`（位于 `slime_plugins/mbridge/glm4moe.py`）在 HuggingFace 与 Megatron 格式之间正确映射普通层和 MTP 层权重。
-
-#### 多机支持
-
-这个示例本身就是多机训练配置。启动前请确认：
-
-- 模型权重和数据集放在所有节点都能访问到的路径；
-- `MASTER_ADDR` 设置为所有节点都能访问到的地址；
-- 在启动 Ray worker 前先取消代理；
-- 提供一个 `HOSTFILE` 列出 worker IP（每行一个），并在启动前 `export HOSTFILE=/path/to/hostfile`；
-- 并行度需要成套调整。默认示例使用 TP=8、PP=4、EP=16、CP=2，rollout 侧则使用 32 张卡 / engine + SGLang DP attention。
-
-如果 rollout GPU 数与 expert 数（160）之间不能整除，可以通过 `--sglang-ep-num-redundant-experts` 增加冗余 expert。
-
-## FP8 Rollout
-
-开源版 GLM-4.7 的 FP8 checkpoint 使用的是 per-channel 量化，目前无法在 SGLang 中直接启用 DeepEP。可以利用 slime 自带工具将其转换为 128x128 的 per-block FP8 checkpoint：
-
-```bash
-cd /root/slime
-python tools/convert_hf_to_fp8.py \
-    --model-dir $BASE_DIR/GLM-4.7-355B-A32B/ \
-    --save-dir $BASE_DIR/GLM-4.7-355B-A32B-FP8/ \
-    --strategy block --block-size 128 128 \
-    --max-workers 4
-```
-
-随后把 `--hf-checkpoint` 改成 `$BASE_DIR/GLM-4.7-355B-A32B-FP8/` 即可开启 FP8 rollout。
-
-一个可参考的 FP8 `SGLANG_ARGS` 配置如下：
-
-```bash
-SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 32
-   --sglang-mem-fraction-static 0.7
-   --sglang-enable-dp-attention
-   --sglang-dp-size 32
-   --sglang-ep-size 32
-   --sglang-moe-dense-tp-size 1
-   --sglang-enable-dp-lm-head
-   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 128)
-
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
-
-   --sglang-moe-a2a-backend deepep
-   --sglang-deepep-mode auto
-)
-```
diff --git a/docs/zh/examples/qwen3-4B.md b/docs/zh/examples/qwen3-4B.md
index 1f3adc9c51..2afc70bfbe 100644
--- a/docs/zh/examples/qwen3-4B.md
+++ b/docs/zh/examples/qwen3-4B.md
@@ -280,10 +280,10 @@ ray job submit ... \
 
 ⚠️  在进行训推分离的时候，每个 sglang server 上的并发度太大，超过了 sglang 默认的 cuda graph 的并发度（默认最大 160），影响推理速度。可以用以下 2 种方式进行调整：
 
-1. 通过 `--rollout-server-concurrency` 限制发给一个 sglang server 的最大并发量，例如：
+1. 通过 `--sglang-server-concurrency` 限制发给一个 sglang server 的最大并发量，例如：
 
    ```bash
-   --rollout-server-concurrency 160
+   --sglang-server-concurrency 160
     ```
 
 2. 使用 `--sglang-cuda-graph-bs`，即 sglang 原生的 `--cuda-graph-bs`, 增大 sglang 初始化的 cuda graph 数量，例如：
diff --git a/docs/zh/get_started/customization.md b/docs/zh/get_started/customization.md
index 509badf3f6..5d1e3ddffd 100644
--- a/docs/zh/get_started/customization.md
+++ b/docs/zh/get_started/customization.md
@@ -28,6 +28,7 @@ slime 通过函数路径参数提供了广泛的自定义能力。这些参数
 | [`--custom-megatron-init-path`](#17-megatron-hook) | Megatron 设置后的自定义初始化。 |
 | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hook) | log probability 计算前的自定义逻辑。 |
 | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hook) | 每个训练步骤前的自定义逻辑。 |
+| [`--slime-router-middleware-paths`](#18-slime-router-中间件---slime-router-middleware-paths) | 向 slime router 添加自定义中间件。 |
 
 ## 详细接口参考
 
@@ -401,14 +402,27 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler
 
 ---
 
-### 18. MoE 路由重放
+### 18. slime Router 中间件 (`--slime-router-middleware-paths`)
+
+**用途**: 向 slime router 添加自定义中间件用于请求处理。
+
+**使用场景**:
+- 请求/响应转换
+- 自定义路由逻辑
+- 缓存和优化
+
+---
+
+### 19. MoE 路由重放
 
 通过记录和重放专家路由决策来稳定 MoE RL 训练。
 
 | 参数 | 说明 |
 | --- | --- |
 | `--use-routing-replay` | 训练中前向-反向路由一致性。([arXiv:2507.18071](https://arxiv.org/abs/2507.18071)) |
-| `--use-rollout-routing-replay` | R3：在训练时重放 rollout 阶段的路由。slime 默认的 `sglang_router` 路径支持该功能。([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) |
+| `--use-rollout-routing-replay` | R3：在训练时重放 rollout 阶段的路由。**需要 `--use-slime-router`**。([arXiv:2510.11370](https://arxiv.org/abs/2510.11370)) |
+
+关于 R3 和 SlimeRouter 的详细说明，请参阅 [Slime Router](../advanced/slime-router.md)。
 
 ## 自定义函数路径的测试
 
diff --git a/docs/zh/get_started/quick_start.md b/docs/zh/get_started/quick_start.md
index 463af729f4..a308ae70ec 100644
--- a/docs/zh/get_started/quick_start.md
+++ b/docs/zh/get_started/quick_start.md
@@ -70,7 +70,7 @@ hf download --repo-type dataset zhuzilin/aime-2024 \
 
 当使用 Megatron 作为训练后端时，需要先将 Hugging Face 格式的模型权重转换为 Megatron `torch_dist` 格式。
 
-首先，加载目标模型的配置文件。`slime/scripts/models` 目录下包含了支持模型的配置文件。需要 `source` 对应模型的脚本，将配置参数加载到当前环境中。此处我们以 GLM4-9B 模型为例子，对于 Qwen3-4B、Qwen3.5、Qwen3.6、GLM-4.7-Flash、Qwen3-30B-A3B，是类似的。
+首先，加载目标模型的配置文件。`slime/scripts/models` 目录下包含了支持模型的配置文件。需要 `source` 对应模型的脚本，将配置参数加载到当前环境中。此处我们以 GLM4-9B 模型为例子，对于 Qwen3-4B，GLM-4.7-Flash，Qwen3-30B-A3B，是类似的。
 
 ```bash
 cd /root/slime
@@ -578,5 +578,5 @@ ray job submit --address="http://127.0.0.1:8265" \
 slime 针对大规模混合专家（MoE）模型的分布式训练进行了深度优化。我们提供了一些端到端的训练案例以供参考：
 
 - [示例：8xH100 训练 GLM-4.7-Flash](../examples/glm4.7-30B-A3B.md)
-- [示例：64xH100 训练 GLM-4.7](../examples/glm4.7-355B-A32B.md)
+- [示例：64xH100 训练 GLM-4.5](../examples/glm4.5-355B-A32B.md)
 - [示例：128xH100 训练 DeepSeek-R1](../examples/deepseek-r1.md)
diff --git a/docs/zh/get_started/usage.md b/docs/zh/get_started/usage.md
index dede294827..1db38f4215 100644
--- a/docs/zh/get_started/usage.md
+++ b/docs/zh/get_started/usage.md
@@ -37,7 +37,8 @@
 
 slime 支持多种训练后端，可以通过 `--train-backend` 参数进行选择：
 
-- `megatron`（默认）：使用 Megatron-LM 作为训练后端，支持大规模模型的高效训练。
+- `megatron`（默认）：使用 Megatron-LM 作为训练后端，支持大规模模型的高效训练；
+- `fsdp`（实验性）：使用 PyTorch FSDP 作为训练后端，可以直接加载 HuggingFace 格式权重，无需转换。
 
 ### 加载 megatron
 
@@ -193,6 +194,7 @@ sglang 的加载非常简单，只需要：
   注意：在策略蒸馏 (OPD) 现在与 advantage estimator 正交，使用 `--use-opd` 和 `--opd-kl-coef` 可以在任意 estimator 之上启用 OPD。
 - `--calculate-per-token-loss`：slime 中默认的方案是 per sample loss，即 `mean(sum(sample_i) / len(sample_i))`，如果需要计算 per token loss，即 `sum(sum(sample_i)) / sum(len(sample_i))`，可以开启 `--calculate-per-token-loss`；
 - `--use-tis`：如果需要开启 tis（https://fengyao.notion.site/off-policy-rl），可以开启这一设置；
+- `--true-on-policy-mode`：开启 True On-Policy 模式，即在训练过程中严格保证数据是当前策略生成的。
 
 #### GRPO 算法
 
@@ -260,24 +262,6 @@ PPO 相关参数：
 - `--value-clip`：value loss 的 clip 范围；
 - `--kl-coef`：KL penalty 系数，用于 reward shaping。
 
-### 高级 Megatron 配置（--megatron-config-path）
-
-对于 PPO 场景，可以使用 `--megatron-config-path` 指定一个 YAML 文件，对 actor / critic 分别覆盖 Megatron 参数。常见用途包括给 critic 设置不同的 `lr`，或者分别指定 `load` / `save` 等路径。
-
-```yaml
-megatron:
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-```
-
-> **注意：** 当前该配置只支持 PPO；并且当前 PPO 下 actor 和 critic 的 Megatron 并行配置必须保持一致。建议把并行相关参数继续写在公共 CLI 中，只把角色差异项放在 YAML 里。详见 [Megatron Config：按角色覆盖训练参数](../advanced/megatron-config.md)。
-
 ## 自定义 rollout 函数
 
 slime 支持不同程度的自定义数据生成（rollout）。
@@ -384,36 +368,6 @@ slime 会用 [sglang-router](https://github.com/sgl-project/sglang/tree/main/sgl
 
 当通过 `--sglang-router-ip` 与 `--sglang-router-port` 来配置传入一个外部的 router，此时 slime 不再会在内部启动一个 router，而是会把所有的 server 都注册在这个外部 router 上。这时可以利用这个外部的 router 地址来实现更复杂的数据生成流程。注意 router 是支持 openai compatible api 的。
 
-### 高级引擎配置（--sglang-config）
-
-对于高级部署场景，可以使用 `--sglang-config` 指定一个 YAML 文件，来配置服务器组、多模型部署以及选择性权重更新。
-
-**多模型部署**允许同时服务多个模型（例如一个接收权重更新的 actor 模型和一个冻结的 reference/reward 模型）：
-
-```yaml
-sglang:
-  - name: actor
-    update_weights: true          # 接收训练的权重更新（默认）
-    server_groups:
-      - worker_type: regular
-        num_gpus: 8
-        num_gpus_per_engine: 4
-  - name: ref
-    model_path: /path/to/ref_model
-    update_weights: false          # 冻结，不更新权重
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 2
-```
-
-每个模型都有自己独立的 router。每个模型的 router 信息可通过 `args.sglang_model_routers`（一个将模型名映射到 `(ip, port)` 元组的字典）访问。自定义 rollout 函数可以使用 `slime.rollout.sglang_rollout` 中的 `get_model_url(args, "ref")` 来将请求路由到指定模型。
-
-**服务器组功能：**
-- `worker_type`：`regular`、`prefill`、`decode` 或 `placeholder`（预留 GPU 位置但不创建引擎）
-- `overrides`：SGLang `ServerArgs` 字段覆盖字典，会叠加在 `--sglang-*` CLI 参数之上
-- `num_gpus_per_engine`：每组的 TP 大小覆盖
-
 ## megatron 使用方法
 
 slime 通过复用 `megatron.training` 目录下的常规函数，如 `parse_args`， `save_checkpoint`，`load_checkpoint`，从而实现对不同版本以及轻度魔改的 megatron 的支持。所以在使用时，需要保证 `PYTHONPATH` 中能访问到 megatron，例如在运行时加入 `export PYTHONPATH=/root/Megatron-LM`。
diff --git a/docs/zh/index.rst b/docs/zh/index.rst
index 36d3d79eb2..7c254194cd 100644
--- a/docs/zh/index.rst
+++ b/docs/zh/index.rst
@@ -34,21 +34,20 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外
 
    examples/glm4.7-30B-A3B.md
    examples/qwen3-30B-A3B.md
-   examples/glm4.7-355B-A32B.md
+   examples/glm4.5-355B-A32B.md
    examples/deepseek-r1.md
 
 .. toctree::
    :maxdepth: 1
    :caption: 高级特性
 
+   advanced/slime-router.md
    advanced/on-policy-distillation.md
    advanced/speculative-decoding.md
    advanced/low-precision.md
    advanced/reproducibility.md
    advanced/fault-tolerance.md
    advanced/pd-disaggregation.md
-   advanced/sglang-config.md
-   advanced/megatron-config.md
    advanced/arch-support-beyond-megatron.md
 
 .. toctree::
@@ -67,7 +66,6 @@ slime 是 GLM-4.7、GLM-4.6、GLM-4.5 背后的 RL 训练框架。除此之外
 
    developer_guide/ci.md
    developer_guide/debug.md
-   developer_guide/trace.md
    developer_guide/profiling.md
 
 .. toctree::
diff --git a/examples/README.md b/examples/README.md
index a1e3bd251e..a64fe249d6 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -6,8 +6,8 @@ These examples provide concrete examples to leverage slime in your own RL workfl
 
 - **[eval_multi_task](./eval_multi_task)**: Example for supporting evaluation multiple tasks with different configs.
 - **[fully_async](./fully_async)**: Demonstrates fully asynchronous rollout generation for higher efficiency.
-- **[geo3k_vlm](./geo3k_vlm)**: Training VLMs on a single-turn reasoning task using GRPO on the GEO3K dataset.
-- **[geo3k_vlm_multi_turn](./geo3k_vlm_multi_turn)**: VLM multi-turn training on Geo3k dataset.
+- **[geo3k_vlm](./geo3k_vlm)**: Training VLMs with FSDP on a single-turn reasoning task using GRPO on the GEO3K dataset.
+- **[geo3k_vlm_multi_turn](./geo3k_vlm_multi_turn)**: VLM multi-turn training (FSDP backend) on Geo3k dataset.
 - **[low_precision](./low_precision)**: Examples of FP8 training and inference for improved throughput and stability.
 - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `slime`.
 - **[on_policy_distillation](./on_policy_distillation)**: Example implementation for on-policy distillation, extending the reinforcement learning pipeline to support teacher–student distillation directly within on-policy training.
@@ -17,3 +17,5 @@ These examples provide concrete examples to leverage slime in your own RL workfl
 - **[strands_sglang](./strands_sglang)**: Integration example with the Strands-Agents scaffolding framework.
 - **[tau-bench](./tau-bench)**: Training in an agentic multi-turn tool use environment (Tau-bench).
 - **[train_infer_mismatch_helper](./train_infer_mismatch_helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS).
+- **[true_on_policy](./true_on_policy)**: Ensures strictly equal log probabilities between inference (SGLang) and training engines.
+- **[true_on_policy_vlm](./true_on_policy_vlm)**: "True On-Policy" training demonstration for VLM (Qwen3-VL).
diff --git a/examples/eval_multi_task/multi_task.yaml b/examples/eval_multi_task/multi_task.yaml
index bad2d61412..83ae67f8ac 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    # hf download --repo-type dataset zyzshishui0627/gpqa_diamond --local-dir /root/gpqa
+    - name: gpqa    # huggingface-cli 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   # hf download --repo-type dataset zyzshishui0627/IFBench --local-dir /root/ifbench
+    - name: ifbench   # huggingface-cli 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/fully_async/fully_async_rollout.py b/examples/fully_async/fully_async_rollout.py
index 1df005ed02..7208365c18 100644
--- a/examples/fully_async/fully_async_rollout.py
+++ b/examples/fully_async/fully_async_rollout.py
@@ -20,7 +20,7 @@ def get_global_worker(args, data_buffer):
     with _worker_lock:
         if _global_worker is None or not _global_worker.worker_thread.is_alive():
             print("Creating new global async worker...")
-            _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.rollout_server_concurrency)
+            _global_worker = AsyncRolloutWorker(args, data_buffer, concurrency=args.sglang_server_concurrency)
             _global_worker.start()
         return _global_worker
 
diff --git a/examples/geo3k_vlm/README.md b/examples/geo3k_vlm/README.md
index ba1ccdbb8d..9841e7641d 100644
--- a/examples/geo3k_vlm/README.md
+++ b/examples/geo3k_vlm/README.md
@@ -1,26 +1,18 @@
-# VLM Single-Turn RL
+# VLM Single-Turn RL (FSDP & Megatron)
 
-Training VLMs with Megatron on single-turn reasoning task using GRPO on the [GEO3K dataset](https://huggingface.co/datasets/hiyouga/geometry3k). We used processed version [here](https://huggingface.co/datasets/chenhegu/geo3k_imgurl).
+Training VLMs with FSDP or Megatron on single-turn reasoning task using GRPO on the [GEO3K dataset](https://huggingface.co/datasets/hiyouga/geometry3k). We used processed version [here](https://huggingface.co/datasets/chenhegu/geo3k_imgurl).
 
 Supported models:
 * Qwen2.5-VL
-* Qwen3-VL (Dense and MoE)
-* Qwen3.5 (Dense and MoE)
+* Qwen3-VL (Dense and Moe)
 
 Note: Please make sure the cudnn version in the environment is 9.16.0.29 to prevent severe performance regression in conv3d in torch 2.9 mentioned in https://github.com/pytorch/pytorch/issues/168167. Otherwise, you can reinstall cudnn with:
 ```bash
 pip install nvidia-cudnn-cu12==9.16.0.29
 ```
 
-**Important:** We use [Megatron Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to support multimodal models. However, not all Megatron arguments are passed through to Megatron Bridge — you may need to set some manually [here](https://github.com/THUDM/slime/blob/de84e10d468dcb726e1199fd6bd16aa9538aed09/slime/backends/megatron_utils/model_provider.py#L89) (currently only parallelization-related arguments are passed). For example, for Qwen3-VL-30B-A3B you may need to add:
-```python
-provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
-provider.freeze_language_model = False
-provider.freeze_vision_model = False
-```
-
 <p align="center">
-  <img src="fsdp_vs_megatron.png" alt="Reward Plot" width="800">
+  <img src="fsdp_vs_megatron.png" alt="FSDP vs Megatron Reward Plot" width="800">
 </p>
 
 ## Data Preparation (For SFT Training)
@@ -63,6 +55,9 @@ export WANDB_API_KEY=your_wandb_api_key
 # Megatron backend (default -> Qwen3-VL-8B-Instruct + Megatron)
 ./examples/geo3k_vlm/run_geo3k_vlm.sh
 
+# FSDP backend
+SLIME_SCRIPT_TRAIN_BACKEND=fsdp ./examples/geo3k_vlm/run_geo3k_vlm.sh
+
 # With different model
 SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.sh
 
@@ -74,6 +69,7 @@ SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.
 
 | Environment Variable | Default | Description |
 |---------------------|---------|-------------|
+| `SLIME_SCRIPT_TRAIN_BACKEND` | `megatron` | Training backend (`megatron` or `fsdp`) |
 | `SLIME_SCRIPT_MODEL_NAME` | `Qwen3-VL-8B-Instruct` | Model name |
 | `SLIME_SCRIPT_DATASET_NAME` | `chenhegu/geo3k_imgurl` | HuggingFace dataset name |
 | `SLIME_SCRIPT_NUM_GPUS` | `8` | Number of GPUs |
@@ -92,11 +88,6 @@ SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-4B-Instruct ./examples/geo3k_vlm/run_geo3k_vlm.
 - `Qwen3-VL-30B-A3B-Thinking`
 - `Qwen3-VL-235B-A22B-Thinking`
 
-#### Qwen3.5 Series
-We provide an [example](./run_geo3k_qwen35.sh) for Qwen3.5-35B-A3B. To support other Qwen3.5 models, add a model config file in `scripts/models/` and update the model name and config path in the script accordingly.
-
-Since Megatron does not currently support packing for GDN, you must set `--qkv-format bshd`, `--micro-batch-size 1`, and remove `--use-dynamic-batch-size`.
-
 ## Notes
 
 ### Reward Model Configuration
diff --git a/examples/geo3k_vlm/run_geo3k_qwen35.sh b/examples/geo3k_vlm/run_geo3k_qwen35.sh
deleted file mode 100644
index 8b402c8dfd..0000000000
--- a/examples/geo3k_vlm/run_geo3k_qwen35.sh
+++ /dev/null
@@ -1,209 +0,0 @@
-#!/bin/bash
-
-# Qwen3.5-35B-A3B VL RL training on geo3k dataset
-
-pip install -U transformers
-
-# IMPORTANT: This branch is specially modified for slime's current Megatron
-# version and Qwen3.5 from the main Megatron Bridge. Other models are not verified!
-# To restore the original Megatron Bridge, run:
-#   pip install git+https://github.com/fzyzcjy/Megatron-Bridge.git@dev_rl --no-build-isolation
-# TODO: Remove this once Megatron & Megatron Bridge are upgraded upstream.
-pip install git+https://github.com/coding-famer/Megatron-Bridge-slime.git@qwen35 --no-build-isolation
-
-# Configuration
-TRAIN_BACKEND="megatron"
-MODEL_NAME="Qwen3_5-35B-A3B"
-DATASET_NAME=${SLIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"}
-NUM_GPUS=${SLIME_SCRIPT_NUM_GPUS:-8}
-DATASET_LOCAL_NAME=$(basename "$DATASET_NAME")
-
-MODEL_NAME_LOWER=$(echo "$MODEL_NAME" | tr '[:upper:]' '[:lower:]')
-
-# External Ray flag
-if [ -z "$SLIME_SCRIPT_EXTERNAL_RAY" ] || [ "$SLIME_SCRIPT_EXTERNAL_RAY" = "0" ]; then
-   USE_EXTERNAL_RAY=0
-else
-   USE_EXTERNAL_RAY=1
-fi
-
-# Cleanup
-pkill -9 sglang
-sleep 3
-if [ "$USE_EXTERNAL_RAY" = "0" ]; then
-   ray stop --force
-   pkill -9 ray
-fi
-pkill -9 slime
-sleep 3
-if [ "$USE_EXTERNAL_RAY" = "0" ]; then
-   pkill -9 ray
-fi
-pkill -9 slime
-pkill -9 redis
-
-set -ex
-
-export PYTHONBUFFERED=16
-
-# Detect NVLink
-NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
-if [ "$NVLINK_COUNT" -gt 0 ]; then
-   HAS_NVLINK=1
-else
-   HAS_NVLINK=0
-fi
-echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
-
-# Download model and dataset
-mkdir -p /root/models /root/datasets
-if [ ! -d "/root/models/${MODEL_NAME}" ]; then
-   hf download Qwen/${MODEL_NAME} --local-dir /root/models/${MODEL_NAME}
-fi
-if [ ! -d "/root/datasets/${DATASET_LOCAL_NAME}" ]; then
-   hf download --repo-type dataset ${DATASET_NAME} --local-dir /root/datasets/${DATASET_LOCAL_NAME}
-fi
-
-# Common args
-CKPT_ARGS=(
-   --hf-checkpoint /root/models/${MODEL_NAME}
-   --load /root/models/${MODEL_NAME}
-   --megatron-to-hf-mode bridge
-)
-
-ROLLOUT_ARGS=(
-   --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train.parquet
-   --input-key problem
-   --label-key answer
-   --apply-chat-template
-   --rollout-shuffle
-   --rm-type deepscaler
-   --num-rollout 3000
-   --rollout-batch-size 64
-   --n-samples-per-prompt 8
-   --rollout-max-response-len 4096
-   --rollout-temperature 0.8
-   --global-batch-size 512
-)
-
-# required for vlm datasets
-MULTIMODAL_KEYS='{"image": "images"}'
-
-EVAL_ARGS=(
-   --eval-interval 20
-   --eval-prompt-data ${DATASET_LOCAL_NAME} /root/datasets/${DATASET_LOCAL_NAME}/test.parquet
-   --n-samples-per-eval-prompt 1
-   --eval-max-response-len 4096
-)
-
-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
-)
-
-OPTIMIZER_ARGS=(
-   --optimizer adam
-   --lr 1e-6
-   --lr-decay-style constant
-   --weight-decay 0.1
-   --adam-beta1 0.9
-   --adam-beta2 0.98
-)
-
-SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 8
-   --sglang-mem-fraction-static 0.7
-   --sglang-ep-size 8
-   --sglang-cuda-graph-bs 1 2 4 8 16 24 32 40 48 56 64 72 80 88 96 104 112 120 128 136 144 152 160 168 176 184 192 200 208 216 224 232 240 248 256
-
-   # MTP speculative decoding
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 2
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 3
-
-   --sglang-max-running-requests 512
-)
-
-# Wandb args (only if WANDB_API_KEY is set)
-if [ -n "$WANDB_API_KEY" ]; then
-   WANDB_ARGS=(
-      --use-wandb
-      --wandb-project slime-geo3k-vlm
-      --wandb-group ${MODEL_NAME_LOWER}-${TRAIN_BACKEND}
-      --wandb-key ${WANDB_API_KEY}
-      --disable-wandb-random-suffix
-   )
-else
-   WANDB_ARGS=()
-fi
-
-MISC_ARGS=(
-   --colocate
-)
-
-# Backend-specific args
-# megatron backend
-BACKEND_ARGS=(
-   --train-backend megatron
-   # Qwen3.5-35B-A3B has num_query_groups = 2
-   --tensor-model-parallel-size 2
-   --sequence-parallel
-   --pipeline-model-parallel-size 1
-   --context-parallel-size 1
-   --expert-model-parallel-size 8
-   --expert-tensor-parallel-size 1
-   --recompute-granularity full
-   --recompute-method uniform
-   --recompute-num-layers 1
-   --attention-dropout 0.0
-   --hidden-dropout 0.0
-   --accumulate-allreduce-grads-in-fp32
-   --attention-softmax-in-fp32
-   --attention-backend flash
-
-   # Packing is not supported for GDN currently
-   --qkv-format bshd
-   --micro-batch-size 1
-)
-
-SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
-source "${SLIME_DIR}/scripts/models/qwen3.5-35B-A3B.sh"
-
-# Start Ray if not using external Ray
-if [ "$USE_EXTERNAL_RAY" = "0" ]; then
-   export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
-   export no_proxy="127.0.0.1,${MASTER_ADDR}"
-   ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-fi
-
-# Build runtime env
-RUNTIME_ENV_JSON="{
-  \"env_vars\": {
-    \"PYTHONPATH\": \"/root/Megatron-LM/\",
-    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
-  }
-}"
-
-ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train.py \
-   --actor-num-nodes 1 \
-   --actor-num-gpus-per-node ${NUM_GPUS} \
-   --multimodal-keys "${MULTIMODAL_KEYS}" \
-   ${MODEL_ARGS[@]} \
-   ${CKPT_ARGS[@]} \
-   ${ROLLOUT_ARGS[@]} \
-   ${EVAL_ARGS[@]} \
-   ${GRPO_ARGS[@]} \
-   ${OPTIMIZER_ARGS[@]} \
-   ${SGLANG_ARGS[@]} \
-   ${WANDB_ARGS[@]} \
-   ${BACKEND_ARGS[@]} \
-   ${MISC_ARGS[@]}
diff --git a/examples/geo3k_vlm/run_geo3k_vlm.sh b/examples/geo3k_vlm/run_geo3k_vlm.sh
index 098a329906..d4e0c89234 100644
--- a/examples/geo3k_vlm/run_geo3k_vlm.sh
+++ b/examples/geo3k_vlm/run_geo3k_vlm.sh
@@ -1,11 +1,13 @@
 #!/bin/bash
 
 # Qwen3 VL RL training on geo3k dataset
+# Supports both megatron and fsdp training backends
 # Usage: 
+#   SLIME_SCRIPT_TRAIN_BACKEND=fsdp ./run_geo3k_vlm.sh
 #   SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-2B-Instruct ./run_geo3k_vlm.sh
 
 # Configuration
-TRAIN_BACKEND="megatron"
+TRAIN_BACKEND=${SLIME_SCRIPT_TRAIN_BACKEND:-"megatron"}
 MODEL_NAME=${SLIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"}
 DATASET_NAME=${SLIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"}
 NUM_GPUS=${SLIME_SCRIPT_NUM_GPUS:-8}
@@ -154,34 +156,46 @@ MISC_ARGS=(
 )
 
 # Backend-specific args
-# megatron backend
-BACKEND_ARGS=(
-   --train-backend megatron
-   --load /root/models/${MODEL_NAME}
-   --tensor-model-parallel-size 4
-   --sequence-parallel
-   --pipeline-model-parallel-size 1
-   --context-parallel-size 1
-   --expert-model-parallel-size 1
-   --expert-tensor-parallel-size 1
-   --recompute-granularity full
-   --recompute-method uniform
-   --recompute-num-layers 1
-   --use-dynamic-batch-size
-   --max-tokens-per-gpu 4096
-   --attention-dropout 0.0
-   --hidden-dropout 0.0
-   --accumulate-allreduce-grads-in-fp32
-   --attention-softmax-in-fp32
-   --attention-backend flash
-   --megatron-to-hf-mode bridge
-)
-
-# get MODEL_ARGS from scripts/models for megatron backend
-SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
-MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
-# VL models require rotary-base 5000000
-MODEL_ARGS_ROTARY_BASE=5000000 source "${SLIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+if [ "$TRAIN_BACKEND" = "fsdp" ]; then
+   BACKEND_ARGS=(
+      --train-backend fsdp
+      --gradient-checkpointing
+      --sglang-attention-backend fa3
+      --attn-implementation flash_attention_3
+      --update-weight-buffer-size 536870912
+   )
+   MODEL_ARGS=()
+else
+   # megatron backend (default)
+   BACKEND_ARGS=(
+      --train-backend megatron
+      --load /root/models/${MODEL_NAME}
+      --tensor-model-parallel-size 4
+      --sequence-parallel
+      --pipeline-model-parallel-size 1
+      --context-parallel-size 1
+      --expert-model-parallel-size 1
+      --expert-tensor-parallel-size 1
+      --recompute-granularity full
+      --recompute-method uniform
+      --recompute-num-layers 1
+      --use-dynamic-batch-size
+      --max-tokens-per-gpu 4096
+      --attention-dropout 0.0
+      --hidden-dropout 0.0
+      --accumulate-allreduce-grads-in-fp32
+      --attention-softmax-in-fp32
+      --attention-backend flash
+      --megatron-to-hf-mode bridge
+   )
+   
+   # get MODEL_ARGS from scripts/models for megatron backend
+   SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
+   MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
+   # VL models require rotary-base 5000000
+   MODEL_ARGS_ROTARY_BASE=5000000 source "${SLIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+   
+fi
 
 # Start Ray if not using external Ray
 if [ "$USE_EXTERNAL_RAY" = "0" ]; then
diff --git a/examples/geo3k_vlm/run_geo3k_vlm_sft.sh b/examples/geo3k_vlm/run_geo3k_vlm_sft.sh
index 06b16d5f92..e2cf44d7cb 100644
--- a/examples/geo3k_vlm/run_geo3k_vlm_sft.sh
+++ b/examples/geo3k_vlm/run_geo3k_vlm_sft.sh
@@ -1,4 +1,4 @@
-TRAIN_BACKEND="megatron"
+TRAIN_BACKEND=${SLIME_SCRIPT_TRAIN_BACKEND:-"megatron"}
 MODEL_NAME=${SLIME_SCRIPT_MODEL_NAME:-"Qwen3-VL-8B-Instruct"}
 DATASET_NAME=${SLIME_SCRIPT_DATASET_NAME:-"chenhegu/geo3k_imgurl"}
 NUM_GPUS=${SLIME_SCRIPT_NUM_GPUS:-8}
@@ -82,6 +82,7 @@ SFT_ARGS=(
    --rollout-function-path slime.rollout.sft_rollout.generate_rollout
    --prompt-data /root/datasets/${DATASET_LOCAL_NAME}/train_formatted.parquet
    --input-key messages
+   --apply-chat-template
    --rollout-shuffle
    --num-epoch 3000
    --rollout-batch-size 128
@@ -121,33 +122,42 @@ else
 fi
 
 # Backend-specific args
-# megatron backend
-BACKEND_ARGS=(
-  --train-backend megatron
-  --tensor-model-parallel-size 4
-  --sequence-parallel
-  --pipeline-model-parallel-size 1
-  --context-parallel-size 1
-  --expert-model-parallel-size 1
-  --expert-tensor-parallel-size 1
-  --recompute-granularity full
-  --recompute-method uniform
-  --recompute-num-layers 1
-  --use-dynamic-batch-size
-  --max-tokens-per-gpu 4096
-  --attention-dropout 0.0
-  --hidden-dropout 0.0
-  --accumulate-allreduce-grads-in-fp32
-  --attention-softmax-in-fp32
-  --attention-backend flash
-  --megatron-to-hf-mode bridge
-)
+if [ "$TRAIN_BACKEND" = "fsdp" ]; then
+    BACKEND_ARGS=(
+      --train-backend fsdp
+      --gradient-checkpointing
+      --attn-implementation flash_attention_3
+      --update-weight-buffer-size 536870912
+    )
+else
+    # megatron backend (default)
+    BACKEND_ARGS=(
+      --train-backend megatron
+      --tensor-model-parallel-size 4
+      --sequence-parallel
+      --pipeline-model-parallel-size 1
+      --context-parallel-size 1
+      --expert-model-parallel-size 1
+      --expert-tensor-parallel-size 1
+      --recompute-granularity full
+      --recompute-method uniform
+      --recompute-num-layers 1
+      --use-dynamic-batch-size
+      --max-tokens-per-gpu 4096
+      --attention-dropout 0.0
+      --hidden-dropout 0.0
+      --accumulate-allreduce-grads-in-fp32
+      --attention-softmax-in-fp32
+      --attention-backend flash
+      --megatron-to-hf-mode bridge
+    )
 
-# get MODEL_ARGS from scripts/models for megatron backend
-SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
-MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
-# VL models require rotary-base 5000000
-MODEL_ARGS_ROTARY_BASE=5000000 source "${SLIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+   # get MODEL_ARGS from scripts/models for megatron backend
+   SLIME_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../.." &>/dev/null && pwd)"
+   MODEL_ARGS_FILE=$(echo "$MODEL_NAME" | sed 's/-Instruct//g; s/-Thinking//g; s/Qwen3-VL-/qwen3-/g; s/-2B/-1.7B/g')
+   # VL models require rotary-base 5000000
+   MODEL_ARGS_ROTARY_BASE=5000000 source "${SLIME_DIR}/scripts/models/${MODEL_ARGS_FILE}.sh"
+fi
 
 # Start Ray if not using external Ray
 if [ "$USE_EXTERNAL_RAY" = "0" ]; then
diff --git a/examples/geo3k_vlm_multi_turn/README.md b/examples/geo3k_vlm_multi_turn/README.md
index 23e88e2025..1947e3cad9 100644
--- a/examples/geo3k_vlm_multi_turn/README.md
+++ b/examples/geo3k_vlm_multi_turn/README.md
@@ -30,6 +30,7 @@ The reward model is the default math RM.
 export WANDB_API_KEY=...
 export SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-2B-Instruct
 export SLIME_SCRIPT_NUM_GPUS=4
+export SLIME_SCRIPT_TRAIN_BACKEND=fsdp
 
 # 2) Download the dataset
 hf download --repo-type dataset VeraIsHere/geo3k_imgurl_processed --local-dir /root/datasets/geo3k_imgurl_processed
diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py
index 8481735cec..bc47d65a0d 100644
--- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py
+++ b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py
@@ -15,6 +15,8 @@
 
 NUM_GPUS = int(os.environ.get("SLIME_SCRIPT_NUM_GPUS", "4"))
 EXTERNAL_RAY = int(os.environ.get("SLIME_SCRIPT_EXTERNAL_RAY", "0"))
+TRAIN_BACKEND = os.environ.get("SLIME_SCRIPT_TRAIN_BACKEND", "fsdp").lower()
+assert TRAIN_BACKEND in {"fsdp", "megatron"}
 
 DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed"
 DATA_ROOT = "/root/datasets/geo3k_imgurl_processed"
@@ -102,7 +104,15 @@ def execute():
         f"--sglang-cuda-graph-bs {' '.join(map(str, [1, 2, 4, 8] + list(range(16, 257, 8))))} "
     )
 
-    backend_args = (
+    fsdp_args = (
+        "--train-backend fsdp "
+        "--gradient-checkpointing "
+        "--sglang-attention-backend fa3 "
+        "--attn-implementation flash_attention_3 "
+        "--update-weight-buffer-size 536870912 "
+    )
+
+    megatron_args = (
         "--train-backend megatron "
         f"--load /root/models/{MODEL_NAME} "
         "--tensor-model-parallel-size 4 "
@@ -124,13 +134,18 @@ def execute():
         "--megatron-to-hf-mode bridge "
     )
 
-    megatron_model_type = get_megatron_model_type(MODEL_NAME)
-    os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
-
     misc_args = (
         "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_GPUS} " f"--rollout-num-gpus {NUM_GPUS} " "--colocate "
     )
 
+    if TRAIN_BACKEND == "megatron":
+        backend_args = megatron_args
+        megatron_model_type = get_megatron_model_type(MODEL_NAME)
+        os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
+    else:
+        backend_args = fsdp_args
+        megatron_model_type = None
+
     train_args = (
         f"{ckpt_args} "
         f"{rollout_args} "
diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py
deleted file mode 100644
index da47268341..0000000000
--- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py
+++ /dev/null
@@ -1,154 +0,0 @@
-import os
-
-from slime.utils.external_utils.command_utils import execute_train_npu
-
-MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")
-assert MODEL_NAME in {
-    "Qwen3-VL-2B-Instruct",
-    "Qwen3-VL-4B-Instruct",
-    "Qwen3-VL-8B-Instruct",
-    "Qwen3-VL-2B-Thinking",
-    "Qwen3-VL-4B-Thinking",
-    "Qwen3-VL-8B-Thinking",
-}
-
-EXTERNAL_RAY = int(os.environ.get("SLIME_SCRIPT_EXTERNAL_RAY", "0"))
-TRAIN_BACKEND = os.environ.get("SLIME_SCRIPT_TRAIN_BACKEND", "fsdp").lower()
-assert TRAIN_BACKEND in {"fsdp", "megatron"}
-
-DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed"
-DATA_ROOT = "/path/to/datasets/geo3k_imgurl_processed"
-TRAIN_DATA_PATH = os.path.join(DATA_ROOT, "train.parquet")
-
-
-def get_megatron_model_type(model_name: str) -> str:
-    model_type = model_name.replace("-Instruct", "").replace("-Thinking", "")
-    model_type = model_type.replace("Qwen3-VL-", "qwen3-")
-    return model_type.replace("-2B", "-1.7B")
-
-
-def execute():
-    ckpt_args = f"--hf-checkpoint /path/to/model/checkpoints/{MODEL_NAME} "
-
-    wandb_args = (
-        (
-            "--use-wandb "
-            "--wandb-project slime-dev "
-            "--wandb-group geo3k_vlm_multi_turn "
-            f"--wandb-key '{wandb_api_key}' "
-        )
-        if (wandb_api_key := os.environ.get("WANDB_API_KEY"))
-        else ""
-    )
-
-    rollout_args = (
-        f"--prompt-data {TRAIN_DATA_PATH} "
-        "--input-key problem "
-        "--label-key answer "
-        '--multimodal-keys \'{"image": "images"}\' '
-        "--rm-type math "
-        "--apply-chat-template "
-        "--custom-generate-function-path examples.geo3k_vlm_multi_turn.rollout.generate "
-        "--custom-config-path examples/geo3k_vlm_multi_turn/geo3k_vlm_multi_turn_config.yaml "
-        "--rollout-shuffle "
-        "--num-rollout 3000 "
-        "--rollout-batch-size 32 "
-        "--n-samples-per-prompt 8 "
-        "--rollout-max-response-len 4096 "
-        "--rollout-temperature 1 "
-        "--global-batch-size 256 "
-    )
-
-    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 "
-        "--use-kl-loss "
-    )
-
-    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 1 "
-        "--sglang-mem-fraction-static 0.6 "
-        f"--sglang-cuda-graph-bs {' '.join(map(str, [4, 8] + list(range(16, 257, 8))))} "
-        "--sglang-device npu "
-        "--sglang-disable-radix-cache "
-        "--sglang-chunked-prefill-size 32768 "
-        "--sglang-max-prefill-tokens 4000 "
-        "--sglang-max-total-tokens 327680 "
-    )
-
-    megatron_args = (
-        "--train-backend megatron "
-        f"--load /path/to/model/checkpoints/{MODEL_NAME} "
-        f"--ref-load /path/to/model/checkpoints/{MODEL_NAME} "
-        "--tensor-model-parallel-size 4 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--recompute-granularity full "
-        "--recompute-method uniform "
-        "--recompute-num-layers 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 16384 "
-        "--balance-data "
-        "--attention-dropout 0.0 "
-        "--hidden-dropout 0.0 "
-        "--accumulate-allreduce-grads-in-fp32 "
-        "--attention-softmax-in-fp32 "
-        "--attention-backend flash "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    misc_args = (
-        "--actor-num-nodes 1 "
-        "--actor-num-gpus-per-node 8 "
-        "--rollout-num-gpus 8 "
-        "--no-gradient-accumulation-fusion "
-        "--use-flash-attn "
-    )
-
-    if TRAIN_BACKEND == "megatron":
-        backend_args = megatron_args
-        megatron_model_type = get_megatron_model_type(MODEL_NAME)
-        os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
-    else:
-        exit()
-
-    train_args = (
-        f"{ckpt_args} "
-        f"{rollout_args} "
-        f"{optimizer_args} "
-        f"{grpo_args} "
-        f"{sglang_args} "
-        f"{backend_args} "
-        f"{misc_args} "
-        f"{wandb_args} "
-    )
-
-    execute_train_npu(
-        train_args=train_args,
-        megatron_model_type=megatron_model_type,
-        extra_env_vars=({"WANDB_API_KEY": os.environ["WANDB_API_KEY"]} if os.environ.get("WANDB_API_KEY") else {}),
-    )
-
-
-if __name__ == "__main__":
-    execute()
diff --git a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py b/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py
deleted file mode 100644
index ca873e9c12..0000000000
--- a/examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py
+++ /dev/null
@@ -1,168 +0,0 @@
-import os
-import tempfile
-
-from slime.utils.external_utils.command_utils import execute_train_npu
-
-MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")
-assert MODEL_NAME in {
-    "Qwen3-VL-2B-Instruct",
-    "Qwen3-VL-4B-Instruct",
-    "Qwen3-VL-8B-Instruct",
-    "Qwen3-VL-2B-Thinking",
-    "Qwen3-VL-4B-Thinking",
-    "Qwen3-VL-8B-Thinking",
-}
-
-EXTERNAL_RAY = int(os.environ.get("SLIME_SCRIPT_EXTERNAL_RAY", "0"))
-TRAIN_BACKEND = os.environ.get("SLIME_SCRIPT_TRAIN_BACKEND", "fsdp").lower()
-assert TRAIN_BACKEND in {"fsdp", "megatron"}
-
-DATASET_NAME = "VeraIsHere/geo3k_imgurl_processed"
-DATA_ROOT = "/path/to/datasets/geo3k_imgurl_processed"
-TRAIN_DATA_PATH = os.path.join(DATA_ROOT, "train.parquet")
-
-
-def get_megatron_model_type(model_name: str) -> str:
-    model_type = model_name.replace("-Instruct", "").replace("-Thinking", "")
-    model_type = model_type.replace("Qwen3-VL-", "qwen3-")
-    return model_type.replace("-2B", "-1.7B")
-
-
-def execute():
-    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
-    megatron_config.write(
-        """
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-"""
-    )
-    megatron_config.close()
-
-    ckpt_args = f"--hf-checkpoint /path/to/model/checkpoints/{MODEL_NAME} "
-
-    wandb_args = (
-        (
-            "--use-wandb "
-            "--wandb-project slime-dev "
-            "--wandb-group geo3k_vlm_multi_turn "
-            f"--wandb-key '{wandb_api_key}' "
-        )
-        if (wandb_api_key := os.environ.get("WANDB_API_KEY"))
-        else ""
-    )
-
-    rollout_args = (
-        f"--prompt-data {TRAIN_DATA_PATH} "
-        "--input-key problem "
-        "--label-key answer "
-        '--multimodal-keys \'{"image": "images"}\' '
-        "--rm-type math "
-        "--apply-chat-template "
-        "--custom-generate-function-path examples.geo3k_vlm_multi_turn.rollout.generate "
-        "--custom-config-path examples/geo3k_vlm_multi_turn/geo3k_vlm_multi_turn_config.yaml "
-        "--rollout-shuffle "
-        "--num-rollout 3000 "
-        "--rollout-batch-size 32 "
-        "--n-samples-per-prompt 8 "
-        "--rollout-max-response-len 4096 "
-        "--rollout-temperature 1 "
-        "--global-batch-size 256 "
-    )
-
-    ppo_args = (
-        "--advantage-estimator ppo "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type k1 "
-        "--kl-coef 0.00 "
-        "--entropy-coef 0.00 "
-        "--eps-clip 4e-4 "
-        "--num-critic-only-steps 1 "
-        "--normalize-advantages "
-    )
-
-    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 1 "
-        "--sglang-mem-fraction-static 0.6 "
-        f"--sglang-cuda-graph-bs {' '.join(map(str, [4, 8] + list(range(16, 257, 8))))} "
-        "--sglang-device npu "
-        "--sglang-disable-radix-cache "
-        "--sglang-chunked-prefill-size 32768 "
-        "--sglang-max-prefill-tokens 4000 "
-        "--sglang-max-total-tokens 327680 "
-    )
-
-    megatron_args = (
-        "--train-backend megatron "
-        f"--load /path/to/model/checkpoints/{MODEL_NAME} "
-        f"--ref-load /path/to/model/checkpoints/{MODEL_NAME} "
-        "--tensor-model-parallel-size 4 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--recompute-granularity full "
-        "--recompute-method uniform "
-        "--recompute-num-layers 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 16384 "
-        "--balance-data "
-        "--attention-dropout 0.0 "
-        "--hidden-dropout 0.0 "
-        "--accumulate-allreduce-grads-in-fp32 "
-        "--attention-softmax-in-fp32 "
-        "--attention-backend flash "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    misc_args = (
-        "--actor-num-nodes 1 "
-        "--actor-num-gpus-per-node 8 "
-        "--rollout-num-gpus 8 "
-        "--no-gradient-accumulation-fusion "
-        "--use-flash-attn "
-    )
-
-    if TRAIN_BACKEND == "megatron":
-        backend_args = megatron_args
-        megatron_model_type = get_megatron_model_type(MODEL_NAME)
-        os.environ["MODEL_ARGS_ROTARY_BASE"] = "5000000"
-    else:
-        exit()
-
-    train_args = (
-        f"--megatron-config-path {megatron_config.name} "
-        f"{ckpt_args} "
-        f"{rollout_args} "
-        f"{optimizer_args} "
-        f"{ppo_args} "
-        f"{sglang_args} "
-        f"{backend_args} "
-        f"{misc_args} "
-        f"{wandb_args} "
-    )
-
-    execute_train_npu(
-        train_args=train_args,
-        megatron_model_type=megatron_model_type,
-        extra_env_vars=({"WANDB_API_KEY": os.environ["WANDB_API_KEY"]} if os.environ.get("WANDB_API_KEY") else {}),
-    )
-
-
-if __name__ == "__main__":
-    execute()
diff --git a/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh b/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh
deleted file mode 100644
index 51ab0b99dd..0000000000
--- a/examples/geo3k_vlm_multi_turn/run_grpo_npu.sh
+++ /dev/null
@@ -1,15 +0,0 @@
-export SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-8B-Instruct
-export SLIME_SCRIPT_TRAIN_BACKEND=megatron
-export PYTORCH_ALLOC_CONF=expandable_segments:True
-export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:/root/sglang/python:$PYTHONPATH"
-export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
-export CUDA_DEVICE_MAX_CONNECTIONS=1
-export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
-export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
-export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
-export HYDRA_FULL_ERROR=1
-export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
-export MASTER_PORT=$(shuf -i 20000-65000 -n 1)  # or any free port
-
-
-python examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_grpo_npu.py
diff --git a/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh b/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh
deleted file mode 100644
index dc1b919a8b..0000000000
--- a/examples/geo3k_vlm_multi_turn/run_ppo_npu.sh
+++ /dev/null
@@ -1,15 +0,0 @@
-export SLIME_SCRIPT_MODEL_NAME=Qwen3-VL-8B-Instruct
-export SLIME_SCRIPT_TRAIN_BACKEND=megatron
-export PYTORCH_ALLOC_CONF=expandable_segments:True
-export PYTHONPATH="/root/Megatron-Bridge/src:/root/Megatron-LM/:/root/sglang/python:$PYTHONPATH"
-export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
-export CUDA_DEVICE_MAX_CONNECTIONS=1
-export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
-export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
-export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
-export HYDRA_FULL_ERROR=1
-export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
-export MASTER_PORT=$(shuf -i 20000-65000 -n 1)  # or any free port
-
-
-python examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn_ppo_npu.py
diff --git a/examples/multi_agent/agent_system.py b/examples/multi_agent/agent_system.py
index 5ccb266079..065f868d51 100644
--- a/examples/multi_agent/agent_system.py
+++ b/examples/multi_agent/agent_system.py
@@ -115,10 +115,10 @@ def __init__(self):
     async def rewrite(self, args, problem_statement, previous_solutions: list[str]) -> str:
         """Generates the rewrited solution."""
 
-        # Build the prompt template dynamically.
+        # 动态生成模板
         template = generate_rewriter_template(len(previous_solutions))
 
-        # Populate the template arguments.
+        # 构建参数字典
         format_params = {"problem_statement": problem_statement}
         for i, solution in enumerate(previous_solutions):
             format_params[f"solution{i+1}"] = solution
@@ -136,10 +136,10 @@ def __init__(self):
     async def select(self, args, problem_statement, candidate_solutions: list[str]) -> str:
         """Generates the rewrited solution."""
 
-        # Build the prompt template dynamically.
+        # 动态生成模板
         template = generate_select_template(len(candidate_solutions))
 
-        # Populate the template arguments.
+        # 构建参数字典
         format_params = {"problem_statement": problem_statement}
         for i, solution in enumerate(candidate_solutions):
             format_params[f"solution{i+1}"] = solution
@@ -170,7 +170,7 @@ async def rewrite_worker(args, previous_solutions, problem_statement, worker_id)
 
 async def solver_worker(args, problem_statement, worker_id):
     """
-    Run a single solver pipeline.
+    单组 solver 流程。
     """
 
     try:
@@ -186,10 +186,10 @@ async def solver_worker(args, problem_statement, worker_id):
 
 async def run_agent_system(args, sample):
     """
-    Run `num_parallel` pipelines concurrently.
+    并发运行 num_parallel 组 pipeline。
     """
 
-    args = deepcopy(args)  # Deep copy args because rollout_with_multi_agents mutates them.
+    args = deepcopy(args)  # 深拷贝 args，因为 args 在 rollout_with_multi_agents 中会被修改
     args.sample = sample
     args.results_dict = {"solver": [], "rewriter": [], "selector": []}
 
@@ -219,7 +219,7 @@ def reward_adjustment(samples, reward_weight):
     ]
     rewrited_solutions_raw = await asyncio.gather(*tasks, return_exceptions=True)
 
-    # Filter out failed tasks and keep only valid rewritten solutions.
+    # 处理异常结果
     rewrited_solutions = []
     for _i, result in enumerate(rewrited_solutions_raw):
         if isinstance(result, str):
@@ -258,8 +258,7 @@ def reward_adjustment(samples, reward_weight):
                     args.results_dict["selector"][0].reward = sample.reward
                     break
 
-    ## If the final answer is correct, add a positive bonus to all rewards.
-    ## Otherwise, add a negative penalty to all rewards.
+    ## 如果最终答案正确，对所有reward添加正向的奖励；如果最终答案不正确，对所有reward添加负向的惩罚。
     if args.results_dict["selector"][0].reward == 1:
         reward_adjustment(args.results_dict["solver"], args.correct_reward_weight)
         reward_adjustment(args.results_dict["rewriter"], args.correct_reward_weight)
diff --git a/examples/multi_agent/prompts.py b/examples/multi_agent/prompts.py
index a79047258f..28ae411f7e 100644
--- a/examples/multi_agent/prompts.py
+++ b/examples/multi_agent/prompts.py
@@ -1,4 +1,4 @@
-## Prompt templates used to generate different prompt variants.
+## 定义了一些 prompts 的模板，用于生成不同的 prompts
 
 
 SOLVER_PROMPT_TEMPLATE = """{problem_statement}"""
diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py
index 339b79dd45..a090af63d7 100644
--- a/examples/retool/generate_with_retool.py
+++ b/examples/retool/generate_with_retool.py
@@ -216,15 +216,6 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     """Custom generation function supporting tool calls"""
     assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment."
 
-    # Retried samples (previously aborted / partial) arrive here with stale
-    # rollout state from the first attempt. Clear it so this generation starts
-    # clean; otherwise the concatenation below appends new tokens to old ones
-    # and downstream `slice_log_prob_with_cp` sees a length mismatch.
-    sample.rollout_log_probs = None
-    sample.response = ""
-    sample.response_length = 0
-    sample.loss_mask = None
-
     state = GenerateState(args)
     url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
 
@@ -238,37 +229,22 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     loss_masks = []
     tool_call_count = 0  # Track actual tool call rounds
 
-    if args.rollout_max_context_len is not None:
-        max_context_length = args.rollout_max_context_len
-    else:
-        max_context_length = args.context_parallel_size * args.max_tokens_per_gpu
-
     for turn in range(TOOL_CONFIGS["max_turns"]):
         # Check if total length exceeds max context length
         total_length = len(prompt_tokens_ids) + len(response_token_ids)
+        if args.rollout_max_context_len is not None:
+            max_context_length = args.rollout_max_context_len
+        else:
+            max_context_length = args.context_parallel_size * args.max_tokens_per_gpu
         if total_length >= max_context_length:
             sample.status = Sample.Status.TRUNCATED
             break
 
-        # Clamp per-turn max_new_tokens to the remaining context budget so a
-        # single turn cannot push total_length past max_context_length. Without
-        # this, a turn can append up to rollout_max_response_len tokens on top
-        # of a total that was just barely under the cap, producing samples
-        # that exceed the training-side max_tokens_per_gpu * cp_size budget
-        # and crash the partition/batch code (asserts or OOMs on an oversized
-        # partition).
-        remaining_budget = max_context_length - total_length
-        per_turn_sampling_params = dict(sampling_params)
-        per_turn_sampling_params["max_new_tokens"] = min(
-            sampling_params.get("max_new_tokens", remaining_budget),
-            remaining_budget,
-        )
-
         # Use token IDs instead of text
         current_token_ids = prompt_tokens_ids + response_token_ids
         payload = {
             "input_ids": current_token_ids,
-            "sampling_params": per_turn_sampling_params,
+            "sampling_params": sampling_params,
             "return_logprob": True,  # Request log probabilities for training
         }
 
@@ -309,14 +285,9 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
             sample.rollout_log_probs += cur_log_probs
 
         else:
-            # sglang returned text but no output_token_logprobs — we cannot
-            # recover per-token logprobs for this turn, which would desync
-            # rollout_log_probs from response_token_ids and blow up
-            # `slice_log_prob_with_cp` downstream. Abort the sample so the
-            # fully_async rollout manager returns the whole group to the
-            # buffer for retry instead of poisoning the trainer.
-            sample.status = Sample.Status.ABORTED
-            return sample
+            cur_response = output["text"]
+            cur_response = postprocess_responses(cur_response)
+            cur_response_token_ids = state.tokenizer(cur_response, add_special_tokens=False)["input_ids"]
 
         response += cur_response
         response_token_ids += cur_response_token_ids
@@ -350,26 +321,6 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
                 sample.rollout_log_probs
             ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(sample.rollout_log_probs)} logps"
 
-        # Tool output is appended verbatim and can push total_length past
-        # max_context_length (the per-turn generation was clamped to the
-        # remaining budget, but tool output is unconstrained). Trim tail
-        # tokens so the final sample fits the training budget exactly.
-        overflow = len(prompt_tokens_ids) + len(response_token_ids) - max_context_length
-        if overflow > 0:
-            response_token_ids = response_token_ids[:-overflow]
-            loss_masks = loss_masks[:-overflow]
-            if sample.rollout_log_probs is not None:
-                sample.rollout_log_probs = sample.rollout_log_probs[:-overflow]
-            # Resync the text field from the trimmed token list so
-            # reward_func's `sample.prompt + sample.response` matches what
-            # the model was actually trained on. decode(tokenize(text)) can
-            # be lossy on some tokenizers (whitespace / special-token
-            # collapse), but reward_func's regex is whitespace-robust and
-            # the trainer sees tokens, not text — so the drift is safe.
-            response = state.tokenizer.decode(response_token_ids)
-            sample.status = Sample.Status.TRUNCATED
-            break
-
         if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]:
             break
 
diff --git a/examples/search-r1/README_zh.md b/examples/search-r1/README_zh.md
index f7d9fbb0d9..a0d273a278 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
-hf download Qwen/Qwen2.5-3B --local-dir /root/Qwen2.5-3B
+huggingface-cli 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 1d5864a47d..0101fc0e69 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
-hf download Qwen/Qwen3-8B --local-dir /root/models/Qwen/Qwen3-8B
+huggingface-cli download Qwen/Qwen3-8B --local-dir /root/models/Qwen/Qwen3-8B
 
 # mcore checkpoint
 cd /root/slime
diff --git a/examples/strands_sglang/generate_with_strands.py b/examples/strands_sglang/generate_with_strands.py
index 42484e1298..efeb13e104 100644
--- a/examples/strands_sglang/generate_with_strands.py
+++ b/examples/strands_sglang/generate_with_strands.py
@@ -1,9 +1,8 @@
-# Updated with strands-sglang 0.3.2
 import logging
 
 from camel.interpreters import SubprocessInterpreter
 from strands import Agent, tool
-from strands_sglang import SGLangModel, ToolLimiter, get_client_from_slime_args
+from strands_sglang import SGLangClient, SGLangModel, ToolLimiter
 from strands_sglang.tool_parsers import HermesToolParser
 
 from slime.rollout.rm_hub.math_dapo_utils import compute_score as math_dapo_compute_score
@@ -24,6 +23,16 @@
 MAX_TOOL_ITERS = 5
 MAX_TOOL_CALLS = None  # No limit
 
+_client_cache: dict[str, SGLangClient] = {}
+
+
+def get_client(args) -> SGLangClient:
+    """Get shared client for connection pooling."""
+    base_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}"
+    if base_url not in _client_cache:
+        _client_cache[base_url] = SGLangClient.from_slime_args(args, timeout=300.0)
+    return _client_cache[base_url]
+
 
 @tool
 def execute_python_code(code: str) -> str:
@@ -46,7 +55,7 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
     state = GenerateState(args)
     model = SGLangModel(
         tokenizer=state.tokenizer,
-        client=get_client_from_slime_args(args, timeout=300.0),
+        client=get_client(args),
         tool_parser=HermesToolParser(),  # tool parsing for wrapped JSON tool calls
         sampling_params=sampling_params,
     )
@@ -60,14 +69,14 @@ async def generate(args, sample: Sample, sampling_params) -> Sample:
         system_prompt=SYSTEM_PROMPT,
     )
 
-    # remember don't set --apply-chat-template in rollout args, it will make user prompt wrapped twice
     prompt = sample.prompt if isinstance(sample.prompt, str) else sample.prompt[0]["content"]
 
     try:
         await agent.invoke_async(prompt)
         sample.status = Sample.Status.COMPLETED
     except Exception as e:
-        # Default all failed rollouts to TRUNCATED; cutomize your logic here if needed
+        # Always use TRUNCATED instead of ABORTED because slime doesn't properly
+        # handle ABORTED samples in reward processing. See: https://github.com/THUDM/slime/issues/200
         sample.status = Sample.Status.TRUNCATED
         logger.warning(f"TRUNCATED: {type(e).__name__}: {e}")
 
diff --git a/examples/tau-bench/README.md b/examples/tau-bench/README.md
index 07724f3bfb..6471aae306 100644
--- a/examples/tau-bench/README.md
+++ b/examples/tau-bench/README.md
@@ -25,11 +25,11 @@ cd /root/slime/examples/tau-bench
 python tau1_mock.py --local_dir /root/tau-bench/
 ```
 
-Initialize the Qwen3-4B-Instruct-2507 model needed for tool use:
+Initialize the Qwen2.5-3B-Instruct model needed for tool use:
 
 ```bash
 # hf checkpoint
-hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507
+huggingface-cli download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507
 
 # mcore checkpoint
 cd /root/slime
diff --git a/examples/tau-bench/run_qwen3_4B.sh b/examples/tau-bench/run_qwen3_4B.sh
index fa36af5526..a821734012 100644
--- a/examples/tau-bench/run_qwen3_4B.sh
+++ b/examples/tau-bench/run_qwen3_4B.sh
@@ -100,7 +100,7 @@ SGLANG_ARGS=(
    --rollout-num-gpus-per-engine 1
    --sglang-mem-fraction-static 0.7
    # If gemini API reports concurrency limit error, set this parameter to reduce the concurrency
-   # --rollout-server-concurrency 32
+   # --sglang-server-concurrency 32
 )
 
 MISC_ARGS=(
diff --git a/examples/train_infer_mismatch_helper/mis.py b/examples/train_infer_mismatch_helper/mis.py
index 3bbf5cf26e..ef46a2e278 100644
--- a/examples/train_infer_mismatch_helper/mis.py
+++ b/examples/train_infer_mismatch_helper/mis.py
@@ -449,3 +449,45 @@ def add_ppl_metrics(
     rho_squared_seq = torch.exp(2.0 * log_ratio_sum_safe)  # (Π ρ_t)²
     chi2_seq = rho_squared_seq - 1.0
     metrics_append(metrics, "chi2_seq", chi2_seq)
+
+
+def compute_mis_weights_fsdp(
+    args,
+    *,
+    pg_loss: torch.Tensor,
+    train_log_probs: list[torch.Tensor],
+    rollout_log_probs: list[torch.Tensor],
+    loss_masks: list[torch.Tensor],
+    **kwargs: Any,
+) -> tuple[torch.Tensor, list[torch.Tensor], dict[str, torch.Tensor]]:
+    """Compute masked importance sampling weights for FSDP. No context parallelism.
+
+    Args:
+        args: Arguments containing MIS settings (use_tis, tis_mode, etc.)
+        pg_loss: Policy gradient loss, flattened tensor [total_tokens]
+        train_log_probs: Training log probs, list of 1D tensors per sequence
+        rollout_log_probs: Rollout log probs, list of 1D tensors per sequence
+        loss_masks: Loss masks, list of 1D tensors per sequence
+        **kwargs: Additional arguments (cp_rank, cp_size, etc.) for compatibility
+
+    Returns:
+        pg_loss: Policy gradient loss with IS weights applied
+        modified_masks: Modified loss masks after rejection sampling
+        mis_metrics: Metrics dict with flattened tensors
+    """
+    is_weights, modified_masks, is_metrics = compute_mis_weights(
+        args=args,
+        train_log_probs=train_log_probs,
+        rollout_log_probs=rollout_log_probs,
+        loss_masks=loss_masks,
+    )
+
+    result_metrics = {}
+    if is_weights is not None:
+        is_weights_flat = torch.cat(is_weights, dim=0)
+        pg_loss = pg_loss * is_weights_flat
+
+    for key, values in is_metrics.items():
+        result_metrics[f"mis_{key}"] = torch.cat(values, dim=0)
+
+    return pg_loss, modified_masks, result_metrics
diff --git a/examples/true_on_policy/README.md b/examples/true_on_policy/README.md
new file mode 100644
index 0000000000..620564d410
--- /dev/null
+++ b/examples/true_on_policy/README.md
@@ -0,0 +1,65 @@
+# True On-Policy between Training and Inference
+
+True on-policy ensures that the log probs generated by inference engine (SGLang) is strictly equal to the one generated by the training Engine. Here's our [blog](https://lmsys.org/blog/2025-12-03-miles-fsdp/) for more details.
+
+## Examples
+
+### Example 1
+
+In this script, we provide a minimal example to use true-on-policy.
+
+```bash
+python examples/true_on_policy/run_simple.py
+```
+
+### Example 2
+
+This script contains more features for various use cases, and one flag is about the true on policy feature.
+
+```bash
+python scripts/run_qwen3_4b.py --train-backend fsdp --true-on-policy
+```
+
+In order to quickly see the curve, you may use `--mode debug_minimal`, which will skip evaluation and run generation with a very short output sequence length. Since true on policy is unrelated to OSL or answer correctness, this can be used for quick experiments.
+
+### Other Cases
+
+In order to support true on policy for other cases, please refer to the flags changed in the examples above.
+
+### What is Expected to Observe
+
+After running the training, you can see in wandb that the metric `train/train_rollout_logprob_abs_diff` should be exactly `0`. This indicates that there is no difference between the log probabilities from the training and the inference. Without the feature enabled, this value should be nonzero.
+
+### Setup & Results
+We fine-tune Qwen3-4B-Base on dapo-math-17k dataset with max_new_tokens = 2048, and evaluate on aime-2024 dataset with max_new_tokens = 8192.
+Global batch size is 64 × 16. Results are summarized below.
+<p align="center"> <img width="360" alt="diff" src="src/train_rollout_abs_diff.png" /> <img width="360" alt="step_time" src="src/step_time.png" /> <img width="360" alt="rollout_time" src="src/rollout_time.png" /> <img width="360" alt="raw_rewards" src="src/raw_reward.png" /> <img width="360" alt="eval" src="src/aime.png" /> </p>
+
+### Observations
+
+Train–inference-diff is strictly 0, verifying full numerical equivalence between training and inference forward passes. Raw rewards perfectly match the baseline, and rollout time shows an acceptable slowdown.
+
+### Reproduction
+
+Detailed reproduction refers to [this](https://gist.github.com/fzyzcjy/46f9fc096258cf6fb4516ad2ffcefa8c).
+
+## How it is Implemented
+
+The core idea is to make each and every operation in training and inference be bitwise equal. The main code is implemented in [#566](https://github.com/THUDM/slime/pull/566) and [SGLang#12058](https://github.com/sgl-project/sglang/pull/12058).
+
+Briefly speaking, we handled the following components to make them aligned:
+
+* Attention: We use the [Flash Attention 3](https://github.com/Dao-AILab/flash-attention) backend for both training and inference, since it achieves bitwise equal between prefill and decode operations.
+* GEMM: We use [DeepGEMM](https://github.com/deepseek-ai/DeepGEMM) for fast matrix multiplication while preserving true-on-policy, thanks to its algorithm to pick things like tensor core instructions ([SGLang#12142](https://github.com/sgl-project/sglang/pull/12142)).
+* Batch invariant kernels: This is a prerequisite for true on-policy, and we use [the ones](https://github.com/thinking-machines-lab/batch_invariant_ops) from the Thinking Machines Lab.
+* Torch compile: We also utilize [`torch.compile`](https://docs.pytorch.org/docs/stable/generated/torch.compile.html) to speed up by avoiding many tiny kernels.
+* We align numeric operation details between the two systems for simplicity, such as op dtype, detailed kernels, etc. Some operations can also be compiled to speedup ([#603](https://github.com/THUDM/slime/pull/603), [SGLang#12161](https://github.com/sgl-project/sglang/pull/12161)).
+
+In order to more easily align the two parts, we use SGLang's [dumper](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/debug_utils/dumper.py) tool for quick comparisons. (Need [#12622](https://github.com/sgl-project/sglang/pull/12622) and [#12623](https://github.com/sgl-project/sglang/pull/12623) for most convenience.)
+
+## Future Works
+
+We will keep maintaining and enhancing this feature. More specifically, we will:
+
+* Further validate it with more experiments, and compare it with algorithmic fixes such as TIS.
+* Potentially scale it to larger scale models if it is proven to be effective.
diff --git a/examples/true_on_policy/run_simple.py b/examples/true_on_policy/run_simple.py
new file mode 100644
index 0000000000..086c8c2432
--- /dev/null
+++ b/examples/true_on_policy/run_simple.py
@@ -0,0 +1,142 @@
+import os
+
+
+import slime.utils.external_utils.command_utils as U
+
+MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-0.6B")
+assert MODEL_NAME in {"Qwen3-0.6B", "Qwen3-4B"}
+
+MODE = os.environ.get("SLIME_SCRIPT_MODE", "normal")
+assert MODE in {"normal", "debug_minimal", "debug_one_sample"}
+
+NUM_GPUS = int(os.environ.get("SLIME_SCRIPT_NUM_GPUS", "1"))
+
+
+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/gsm8k")
+
+
+def execute():
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} "
+
+    rollout_args = (
+        "--prompt-data /root/datasets/gsm8k/train.parquet "
+        "--input-key messages "
+        "--label-key label "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type math "
+        f"--num-rollout {2 if MODE == 'debug_one_sample' else 3000} "
+        f"--rollout-batch-size {1 if MODE == 'debug_one_sample' else 32} "
+        f"--n-samples-per-prompt {1 if MODE == 'debug_one_sample' else 8} "
+        f"--rollout-max-response-len {2 if MODE == 'debug_one_sample' else 1024} "
+        "--rollout-temperature 1 "
+        # temp remove this to make test easier
+        # "--over-sampling-batch-size 64 "
+        # "--dynamic-sampling-filter-path slime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std "
+        f"--global-batch-size {1 if MODE == 'debug_one_sample' else 256} "
+    )
+
+    eval_args = ""
+    if MODE == "normal":
+        eval_args = (
+            "--eval-interval 20 "
+            "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet "
+            "--n-samples-per-eval-prompt 1 "
+            "--eval-max-response-len 1024 "
+            "--eval-top-k 1 "
+        )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        # "--use-kl-loss "
+        "--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 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--sglang-decode-log-interval 1000 "
+        "--sglang-enable-metrics "
+        f"--sglang-mem-fraction-static {0.2 if MODEL_NAME == 'Qwen3-4B' else 0.4} "
+        f"{'--sglang-disable-cuda-graph ' if MODE == 'debug_one_sample' else ''}"
+    )
+
+    fsdp_args = (
+        # Set to true for FULL_STATE_DICT mode, false for SHARDED_STATE_DICT mode (default)
+        # "--fsdp-full-params "  # Uncomment this line to enable full params mode
+        # Set the bucket size for weight update
+        "--update-weight-buffer-size 536870912 "  # 512MB
+    )
+
+    ci_args = "--ci-test " "--ci-disable-kl-checker "
+
+    misc_args = "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_GPUS} " "--colocate " "--train-backend fsdp "
+
+    if MODEL_NAME == "Qwen3-4B":
+        misc_args += (
+            "--use-dynamic-batch-size "
+            # TODO pick a good value
+            "--max-tokens-per-gpu 2048 "
+        )
+
+    true_on_policy_args = (
+        "--sglang-enable-deterministic-inference "
+        "--sglang-rl-on-policy-target fsdp "
+        "--sglang-attention-backend fa3 "
+        "--attn-implementation flash_attention_3 "
+        "--deterministic-mode "
+        "--true-on-policy-mode "
+    )
+    true_on_policy_envs = {
+        # TODO note: "Ring" in original RL PR, "allreduce:tree" in SGLang
+        # "NCCL_ALGO": "Ring",
+        "NCCL_ALGO": "allreduce:tree",
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
+        "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+    }
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{sglang_args} "
+        f"{U.get_default_wandb_args(__file__)} "
+        f"{eval_args} "
+        f"{fsdp_args} "
+        f"{ci_args} "
+        f"{misc_args} "
+        f"{true_on_policy_args} "
+    )
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=NUM_GPUS,
+        megatron_model_type=None,
+        extra_env_vars={
+            **true_on_policy_envs,
+            "SGLANG_DUMPER_ENABLE": "1" if MODE == "debug_one_sample" else "0",
+            "SGLANG_TEMP_UTILS_ENABLE_DEBUG_PRINT": "1" if MODE == "debug_one_sample" else "0",
+        },
+    )
+
+
+if __name__ == "__main__":
+    prepare()
+    execute()
diff --git a/examples/true_on_policy/src/aime.png b/examples/true_on_policy/src/aime.png
new file mode 100644
index 0000000000..31a04f9275
Binary files /dev/null and b/examples/true_on_policy/src/aime.png differ
diff --git a/examples/true_on_policy/src/raw_reward.png b/examples/true_on_policy/src/raw_reward.png
new file mode 100644
index 0000000000..eaf013f91a
Binary files /dev/null and b/examples/true_on_policy/src/raw_reward.png differ
diff --git a/examples/true_on_policy/src/rollout_time.png b/examples/true_on_policy/src/rollout_time.png
new file mode 100644
index 0000000000..bb4fe9c23c
Binary files /dev/null and b/examples/true_on_policy/src/rollout_time.png differ
diff --git a/examples/true_on_policy/src/step_time.png b/examples/true_on_policy/src/step_time.png
new file mode 100644
index 0000000000..469706dc19
Binary files /dev/null and b/examples/true_on_policy/src/step_time.png differ
diff --git a/examples/true_on_policy/src/train_rollout_abs_diff.png b/examples/true_on_policy/src/train_rollout_abs_diff.png
new file mode 100644
index 0000000000..b5bac13c7a
Binary files /dev/null and b/examples/true_on_policy/src/train_rollout_abs_diff.png differ
diff --git a/examples/true_on_policy_vlm/README.md b/examples/true_on_policy_vlm/README.md
new file mode 100644
index 0000000000..5cf52bfadb
--- /dev/null
+++ b/examples/true_on_policy_vlm/README.md
@@ -0,0 +1,23 @@
+# True On-Policy between Training and Inference for VLM
+
+This example demonstrates true on-policy training with Qwen3-VL dense model on FSDP. The core concepts and expected observations are the same as [true_on_policy](../true_on_policy/README.md).
+
+<p align="center">
+  <img src="diff.png" alt="Training Inference Log Prob Diff" width="800">
+</p>
+
+## Usage
+
+```bash
+SLIME_SCRIPT_NUM_GPUS=8 python examples/true_on_policy_vlm/run_simple.py
+```
+
+## How it is Implemented
+
+For the text backbone, please refer to [true_on_policy for the text-only model](../true_on_policy/README.md).
+
+For the VLM, we only need to ensure that the image encoder behaves as expected. Please refer to [SGLang#14636](https://github.com/sgl-project/sglang/pull/14636). We need to align numeric operation details between the two systems, so that the ViT forward pass matches the behavior in both SGLang and transformers.
+
+## Notes
+
+It is expected that the true-on-policy version is slower.
\ No newline at end of file
diff --git a/examples/true_on_policy_vlm/diff.png b/examples/true_on_policy_vlm/diff.png
new file mode 100644
index 0000000000..ab4d68030f
Binary files /dev/null and b/examples/true_on_policy_vlm/diff.png differ
diff --git a/examples/true_on_policy_vlm/run_simple.py b/examples/true_on_policy_vlm/run_simple.py
new file mode 100644
index 0000000000..9d866c8361
--- /dev/null
+++ b/examples/true_on_policy_vlm/run_simple.py
@@ -0,0 +1,137 @@
+import os
+
+import slime.utils.misc as U
+from slime.utils.external_utils.command_utils import execute_train, get_default_wandb_args
+
+MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")
+assert MODEL_NAME in {"Qwen2.5-VL-3B-Instruct", "Qwen3-VL-2B-Instruct", "Qwen3-VL-4B-Instruct", "Qwen3-VL-8B-Instruct"}
+
+NUM_GPUS = int(os.environ.get("SLIME_SCRIPT_NUM_GPUS", "1"))
+EXTERNAL_RAY = int(os.environ.get("SLIME_SCRIPT_EXTERNAL_RAY", "0"))
+
+
+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}")
+    dataset_name = "chenhegu/geo3k_imgurl"
+    _, partial_name = dataset_name.split("/")
+    U.exec_command(f"hf download --repo-type dataset {dataset_name} --local-dir /root/datasets/{partial_name}")
+
+
+def execute():
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} "
+
+    rollout_args = (
+        "--prompt-data /root/datasets/geo3k_imgurl/train.parquet "
+        "--input-key problem "
+        "--label-key answer "
+        '--multimodal-keys \'{"image": "images"}\' '
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type math "
+        "--num-rollout 3000 "
+        "--rollout-batch-size 64 "
+        "--n-samples-per-prompt 8 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 512 "
+    )
+
+    eval_args = (
+        "--eval-interval 20 "
+        "--eval-prompt-data geo3k /root/datasets/geo3k_imgurl/test.parquet "
+        "--n-samples-per-eval-prompt 1 "
+        "--eval-max-response-len 4096 "
+        "--eval-top-k 1 "
+    )
+
+    grpo_args = (
+        "--advantage-estimator grpo "
+        # "--use-kl-loss "
+        "--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 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--sglang-mem-fraction-static 0.6 "
+        f"--sglang-cuda-graph-bs {' '.join(map(str, [1, 2, 4, 8] + list(range(16, 257, 8))))} "
+    )
+
+    fsdp_args = (
+        # Set to true for FULL_STATE_DICT mode, false for SHARDED_STATE_DICT mode (default)
+        # "--fsdp-full-params "  # Uncomment this line to enable full params mode
+        # Set the bucket size for weight update
+        "--update-weight-buffer-size 536870912 "  # 512MB
+        "--train-backend fsdp "
+        "--gradient-checkpointing "
+        "--sglang-attention-backend fa3 "
+        "--attn-implementation flash_attention_3 "
+    )
+
+    ci_args = "--ci-test " "--ci-disable-kl-checker "
+
+    misc_args = "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_GPUS} " "--colocate "
+
+    # misc_args += (
+    #     "--use-dynamic-batch-size "
+    #     # TODO pick a good value
+    #     "--max-tokens-per-gpu 2048 "
+    # )
+
+    true_on_policy_args = (
+        "--sglang-enable-deterministic-inference "
+        "--sglang-rl-on-policy-target fsdp "
+        "--deterministic-mode "
+        "--true-on-policy-mode "
+    )
+    true_on_policy_envs = {
+        # TODO note: "Ring" in original RL PR, "allreduce:tree" in SGLang
+        # "NCCL_ALGO": "Ring",
+        "NCCL_ALGO": "allreduce:tree",
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
+        "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+        "SGLANG_VLM_CACHE_SIZE_MB": "0",
+    }
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{sglang_args} "
+        f"{fsdp_args} "
+        f"{ci_args} "
+        f"{eval_args} "
+        f"{misc_args} "
+        f"{get_default_wandb_args(__file__)} "
+        f"{true_on_policy_args} "
+    )
+
+    # Submit Ray job
+    execute_train(
+        train_args=train_args,
+        num_gpus_per_node=NUM_GPUS,
+        megatron_model_type=None,
+        extra_env_vars={
+            **true_on_policy_envs,
+        },
+    )
+
+
+if __name__ == "__main__":
+    prepare()
+    execute()
diff --git a/goal_plan.md b/goal_plan.md
deleted file mode 100644
index f7b8827c8b..0000000000
--- a/goal_plan.md
+++ /dev/null
@@ -1,33 +0,0 @@
-### 阶段一：打通Qwen2.5-0.5B GRPO 8卡同步/异步训练（train.py和train_async.py），GSM8K 数据集，loss/reward 收敛与 SGLang backend 基本一致，且满足确定性计算，多次重复运行Loss曲线完全一致。
-
-First Design and RFC by 03/06 
-
-#### 初步方案：
-- 对标SGLang，Slime 在 Ray 内管理 vLLM 的完整生命周期，包括进程拉起、权重同步、推理暂停/恢复
-- 暂不使用Router，SGLang Model Gateway仅只支持SGLang Worker，SlimeRouter仅在 R3 / radix-tree caching 时需要，Qwen2.5-0.5B 非 MoE 且用 token-in/token-out
-- 单vLLM实例，无router，通过vLLMClient 直连本地 vLLM 进程端口
-- 先支持训推不共卡(non-colocate)，权重同步采用NCCL broadcast，对标SGLang update_weights_from_distributed  (默认）
-- 再支持和验证colocate，权重同步采用GPU IPC（vLLM update_weights_from_ipc, update_weights_from_tensor），对标SGLang update_weights_from_tensor，以验证Reproductivity。**IPC 依赖vllm 0.17**
-
-#### 风险：
-- slime, sglang版本依赖，和vllm 0.16的版本依赖冲突(numpy, torch, transformers, etc)
-- slime代码较挫，可靠性差，强依赖preset docker
-- 算力
-
-
-#### Reference
-
-https://thudm.github.io/slime/advanced/reproducibility.html
-
-
-### 阶段二：接入vllm-project/router，支持多实例vLLM
-
-- vllm router forked from SGLang Model Gateway
-
-### 阶段三：多节点大规模验证，MoE模型，optional：验证MTP Speculative Decoding，FP8 rollout 等高级特性
-
-- Model: Qwen/Qwen3-30B-A3B or GLM4.7
-- Parallel: 16卡 or 128卡, Train mixed EP+FSDP, Rollout EP+DP
-- Verify more features:
-    - Bf16 train, FP8 rollout 
-    - MTP Speculative Decoding
diff --git a/rfc-rollout-backend-separation-plan.md b/rfc-rollout-backend-separation-plan.md
deleted file mode 100644
index ef715c2a6d..0000000000
--- a/rfc-rollout-backend-separation-plan.md
+++ /dev/null
@@ -1,271 +0,0 @@
-# RFC: Rollout Separation Plan (EngineGroup Generalization + Executor Cleanup)
-
-## 1. Summary
-
-This RFC proposes backend separation with minimal churn in runtime orchestration.
-
-Scope is four items:
-
-1. Refactor [slime/ray/rollout.py](slime/ray/rollout.py) by generalizing `EngineGroup.start_engines()` and abstracting engine/server creation (no new runtime manager class hierarchy).
-2. Refactor [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) in-place: extract the one SGLang-specific code path (RadixTree) into a strategy hook, rename SGLang-prefixed args to generic names. No class hierarchy, no new files.
-3. Refactor [slime/utils/arguments.py](slime/utils/arguments.py) into shared args + backend arg groups/finalizers.
-4. Decouple [slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) from SGLang internals so FSDP weight sync works with both SGLang and vLLM engines.
-
-## 2. Already Done (Reuse, Do Not Rewrite)
-
-- Unified rollout contracts in [slime/rollout/base_types.py](slime/rollout/base_types.py).
-- Backend client abstraction in [slime/rollout/backends/base_client.py](slime/rollout/backends/base_client.py).
-- Backend adapters in [slime/rollout/backends/sglang_client.py](slime/rollout/backends/sglang_client.py) and [slime/rollout/backends/vllm_client.py](slime/rollout/backends/vllm_client.py).
-- Managed vLLM engine actor in [slime/backends/vllm_utils/vllm_engine.py](slime/backends/vllm_utils/vllm_engine.py).
-- vLLM translation sidecar in [slime/backends/vllm_utils/vllm_translation_sidecar.py](slime/backends/vllm_utils/vllm_translation_sidecar.py).
-
-## 3. Problem
-
-- [slime/ray/rollout.py](slime/ray/rollout.py) mixes shared and backend-specific engine creation paths.
-- [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) is ~95% backend-agnostic but has one inlined SGLang-specific code path (RadixTree, ~14 lines) and uses SGLang-prefixed arg names for generic rollout concepts.
-- [slime/utils/arguments.py](slime/utils/arguments.py) still has SGLang alias behavior in vLLM path.
-- [slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) hard-imports SGLang internals (`FlattenedTensorBucket`, `MultiprocessingSerializer`, `monkey_patch_torch_reductions`) and calls SGLang-specific engine RPC names (`update_weights_from_tensor`, `update_weights_from_distributed`), making FSDP weight sync unusable with vLLM engines.
-
-## 4. Goals and Non-Goals
-
-### Goals
-
-- Keep runtime refactor minimal and localized to `EngineGroup` + creation abstraction.
-- Isolate the one SGLang-specific executor code path behind a strategy hook; keep functions as functions.
-- Rename SGLang-prefixed arg names to generic rollout names to eliminate naming coupling.
-- Reduce backend leakage in argument finalization.
-- Preserve current external behavior, call sites, and import paths.
-
-### Non-Goals
-
-- No rewrite of `SGLangEngine` or `VLLMEngine` internals.
-- No algorithmic changes to GRPO/PPO.
-- No mandatory feature parity for unsupported backend capabilities.
-
-## 5. Design
-
-### 5.1 Runtime: Generalize `EngineGroup` (No New Runtime Manager Classes)
-
-Keep [slime/ray/rollout.py](slime/ray/rollout.py) as the orchestration entry file.
-
-Refactor focus:
-
-1. Generalize `EngineGroup.start_engines()` to call backend-aware creation hooks.
-2. Abstract engine creation and rollout-server assembly helpers.
-3. Keep existing startup function API (`start_rollout_servers`) and return shape.
-
-Proposed helper abstraction points:
-
-- `create_engine_actor_cls(args, worker_type)`
-  - returns `ray.remote(SGLangEngine)` or `ray.remote(VLLMEngine)`.
-- `create_engine_remote(args, actor_cls, scheduling_strategy, ...)`
-  - encapsulates `.options(...).remote(...)` with backend-specific init kwargs.
-- `build_rollout_server(...)`
-  - standardizes `RolloutServer` construction from engine groups.
-
-`EngineGroup` and `RolloutServer` remain the main shared dataclasses.
-
-### 5.2 Executor: Isolate Backend Logic In-Place (No Class Hierarchy)
-
-#### Current state analysis
-
-[slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) (577 lines) is **already ~95% backend-agnostic**:
-
-| Function / Class | Lines | Backend-specific? | Notes |
-|---|---|---|---|
-| `_get_backend_client()` | 6 | Factory only | Delegates to existing `RolloutBackendClient` subclasses |
-| `_apply_backend_response()` | 20 | No | Uses `RolloutBackendResponse` contract |
-| `GenerateState` | 37 | **Naming only** | References `sglang_server_concurrency`, `sglang_dp_size`, `sglang_enable_deterministic_inference` — all are generic rollout concepts with SGLang-prefixed names |
-| `generate()` | 58 | **14 lines** | RadixTree middleware path (L170-183) is 100% SGLang-specific; the else branch (L185-195) already uses `RolloutBackendClient` |
-| `generate_and_rm()` | 60 | No | Shared orchestration (semaphore, custom func, reward) |
-| `generate_and_rm_group()` | 37 | No | Group parallelism + deterministic seeds |
-| `abort()` | 38 | No | Already uses `backend.abort()` |
-| `generate_rollout_async()` | 71 | No | Main loop, filtering, metrics |
-| `eval_rollout()` / `eval_rollout_single_dataset()` | 118 | **Naming only** | `sglang_enable_deterministic_inference` reference |
-| `generate_rollout()` | 41 | No | Sync entry point |
-
-**Conclusion**: the actual backend logic that needs isolation is **one code path** (~14 lines) inside `generate()`. Everything else is either already abstracted through `RolloutBackendClient` or is a naming-only coupling (SGLang-prefixed arg names for generic concepts).
-
-#### Approach
-
-1. **Extract the RadixTree path into a strategy hook** that `generate()` calls conditionally.
-2. **Rename SGLang-prefixed args** to generic names (coordinated with Phase 3 args refactor).
-3. **Keep functions as functions** — they compose well and callers (`train.py`, OPD, multi-agent) import them directly.
-
-#### Concrete changes
-
-**Step 1 — Extract RadixTree strategy from `generate()`**
-
-Current `generate()` has an `if use_radix: ... else: backend.generate(...)` branch.
-Refactor into:
-
-```python
-# slime/rollout/sglang_rollout.py — generate() simplified
-
-async def generate(args, sample, sampling_params):
-    ...
-    input_ids = ...  # shared prompt encoding (unchanged)
-
-    strategy = _get_generate_strategy(args)
-    resp = await strategy(args, sample, input_ids, sampling_params)
-    _apply_backend_response(sample, resp, args)
-    return sample
-```
-
-Two strategies:
-
-```python
-# Still in sglang_rollout.py (no new file needed)
-
-def _get_generate_strategy(args):
-    """Return the generate coroutine to use."""
-    if _is_radix_tree_enabled(args):
-        return _generate_radix_tree      # SGLang-only path
-    return _generate_via_backend_client   # Generic path (SGLang or vLLM)
-
-def _is_radix_tree_enabled(args) -> bool:
-    return (
-        args.use_slime_router
-        and "RadixTreeMiddleware" in getattr(args, "slime_router_middleware_paths", [])
-    )
-
-async def _generate_radix_tree(args, sample, input_ids, sampling_params) -> RolloutBackendResponse:
-    """SGLang RadixTree middleware path — returns normalized response."""
-    from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree
-    url = f"http://{args.rollout_router_ip}:{args.rollout_router_port}/generate"
-    payload = { ... }  # existing payload construction
-    output = await post(url, payload, headers=headers)
-    sample = await postprocess_sample_with_radix_tree(args, sample, output)
-    return _extract_response_from_sample(sample)  # normalize to RolloutBackendResponse
-
-async def _generate_via_backend_client(args, sample, input_ids, sampling_params) -> RolloutBackendResponse:
-    """Generic backend client path — works for SGLang and vLLM."""
-    backend = _get_backend_client(args)
-    base_url = f"http://{args.rollout_router_ip}:{args.rollout_router_port}"
-    req = RolloutBackendRequest(...)
-    return await backend.generate(req, base_url, headers=headers)
-```
-
-**Step 2 — Rename SGLang-prefixed args to generic names**
-
-| Current name | New name | Reason |
-|---|---|---|
-| `sglang_server_concurrency` | `rollout_concurrency` | Controls request parallelism for any backend |
-| `sglang_dp_size` | `rollout_dp_size` | Data-parallel sharding, not SGLang-specific |
-| `sglang_router_ip` / `sglang_router_port` | `rollout_router_ip` / `rollout_router_port` | Router endpoint, backend-agnostic |
-| `sglang_router_policy` | `rollout_router_policy` | Routing strategy |
-| `sglang_enable_deterministic_inference` | `rollout_deterministic_inference` | Seed-based determinism |
-| `vllm_base_url` | (remove) | Folded into `rollout_router_ip:port`, no special case |
-
-Legacy aliases kept in [slime/utils/arguments.py](slime/utils/arguments.py) for one release cycle (coordinated with Phase 3).
-
-**Step 3 — No new files**
-
-The file stays as [slime/rollout/sglang_rollout.py](slime/rollout/sglang_rollout.py) during this phase.
-Optionally rename to `slime/rollout/rollout.py` in Phase 4 cleanup, since the file is backend-agnostic after the refactor.
-
-### 5.3 Arguments/Config Refactor
-
-In [slime/utils/arguments.py](slime/utils/arguments.py), split into:
-
-1. Shared rollout args.
-2. SGLang backend args/validation.
-3. vLLM backend args/validation.
-
-Add backend finalizers:
-
-- `finalize_sglang_args(args)`
-- `finalize_vllm_args(args)`
-
-Move SGLang alias fallback out of shared finalize flow.
-
-### 5.4 Weight Sync: Decouple FSDP `update_weight_utils.py` from SGLang Internals
-
-#### Current state analysis
-
-[slime/backends/fsdp_utils/update_weight_utils.py](slime/backends/fsdp_utils/update_weight_utils.py) (287 lines) has two concrete classes:
-
-| Class | Weight-push method | SGLang coupling |
-|---|---|---|
-| `UpdateWeightFromTensor` | IPC via Gloo gather → `engine.update_weights_from_tensor.remote()` | Imports `FlattenedTensorBucket`, `MultiprocessingSerializer`, `monkey_patch_torch_reductions` directly from `sglang.srt.*` |
-| `UpdateWeightFromDistributed` | NCCL broadcast → `engine.update_weights_from_distributed.remote()` | Calls `engine.init_weights_update_group.remote()` — SGLang engine API |
-
-The abstract base `UpdateWeight` itself is clean (only PyTorch + Ray).
-
-The Megatron side already solved this: [slime/backends/megatron_utils/sglang.py](slime/backends/megatron_utils/sglang.py) centralizes all SGLang imports into one shim. The FSDP side duplicates these imports inline.
-
-#### Coupling points
-
-1. **SGLang utility imports (L13-26)** — `monkey_patch_torch_reductions`, `MultiprocessingSerializer`, `FlattenedTensorBucket` are imported directly from `sglang.srt.*` with try/except version fallbacks.
-2. **`UpdateWeightFromTensor.update_bucket_weights()`** — uses `FlattenedTensorBucket` to flatten tensors, `MultiprocessingSerializer` to serialize, then calls `engine.update_weights_from_tensor.remote()`.
-3. **`UpdateWeightFromDistributed.connect_rollout_engines()`** — calls `engine.init_weights_update_group.remote()` which is an SGLang engine method.
-4. **`UpdateWeightFromDistributed.update_bucket_weights()`** — calls `engine.update_weights_from_distributed.remote()` which is an SGLang engine method.
-
-All four points assume the engine actor exposes SGLang's RPC interface. vLLM engines expose different method names.
-
-#### Approach
-
-**Step 1 — Centralize SGLang imports via a shim (same pattern as Megatron)**
-
-Reuse or mirror the existing [slime/backends/megatron_utils/sglang.py](slime/backends/megatron_utils/sglang.py) pattern:
-
-```python
-# slime/backends/fsdp_utils/sglang_compat.py (new, ~15 lines)
-try:
-    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
-except ImportError:
-    from sglang.srt.patch_torch import monkey_patch_torch_reductions
-
-from sglang.srt.utils import MultiprocessingSerializer
-
-try:
-    from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket
-except ImportError:
-    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket
-```
-
-Then `update_weight_utils.py` imports from `sglang_compat` — one import line instead of five, and only loaded when SGLang is the active backend.
-
-**Step 2 — Abstract engine RPC calls behind a protocol**
-
-The engine actors (`SGLangEngine`, `VLLMEngine`) already expose weight-sync methods but with different names/signatures. Introduce a lightweight protocol or adapter:
-
-```python
-# In update_weight_utils.py or a small helper
-
-def _call_engine_update_tensor(engine, backend: str, **kwargs):
-    """Dispatch IPC weight update to the correct engine method."""
-    if backend == "vllm":
-        return engine.update_weights_from_tensor.remote(**kwargs)  # vLLM uses same name via NcclBridge
-    return engine.update_weights_from_tensor.remote(**kwargs)  # SGLang native
-
-def _call_engine_update_distributed(engine, backend: str, **kwargs):
-    if backend == "vllm":
-        return engine.update_weights_from_distributed.remote(**kwargs)
-    return engine.update_weights_from_distributed.remote(**kwargs)
-
-def _call_engine_init_weight_group(engine, backend: str, **kwargs):
-    if backend == "vllm":
-        return engine.init_weights_update_group.remote(**kwargs)
-    return engine.init_weights_update_group.remote(**kwargs)
-```
-
-> Note: Currently `VLLMEngine` already mirrors these method names (it wraps them via `NcclBridge`), so the dispatch functions may initially be identical. The value is making the indirection explicit so future method-name divergence is handled in one place.
-
-**Step 3 — Lazy-import SGLang utilities only when backend is SGLang**
-
-Move the `FlattenedTensorBucket` / `MultiprocessingSerializer` imports inside `UpdateWeightFromTensor.update_bucket_weights()` behind a lazy import, so the module can be loaded in a vLLM-only environment without SGLang installed.
-
-#### What changes
-
-| File | Change |
-|---|---|
-| `slime/backends/fsdp_utils/sglang_compat.py` | New shim file (~15 lines) centralizing SGLang imports |
-| `slime/backends/fsdp_utils/update_weight_utils.py` | Replace 5 inline SGLang imports with one `from .sglang_compat import ...`; add backend-aware engine dispatch helpers |
-
-#### What does NOT change
-
-- `UpdateWeight` abstract base class — already clean.
-- `UpdateWeightFromDistributed` NCCL logic — the broadcast itself is pure PyTorch; only the engine RPC dispatch gets a thin wrapper.
-- Megatron-side weight sync — already has its own shim, not touched.
-
diff --git a/rfc-vllm-rollout-backend-en.md b/rfc-vllm-rollout-backend-en.md
deleted file mode 100644
index dad61a0731..0000000000
--- a/rfc-vllm-rollout-backend-en.md
+++ /dev/null
@@ -1,401 +0,0 @@
-# RFC: Supporting vLLM as a Rollout Backend in Slime
-
-- **Author**: \<your name\>
-- **Status**: Phase 1 Done
-- **Audience**: Slime rollout/runtime maintainers, RL training maintainers
-- **Last Updated**: 2026-03-03
-
-## 1. Summary
-
-This RFC proposes adding **vLLM** as a first-class rollout backend in Slime while keeping the existing SGLang behavior unchanged and the GRPO training pipeline unaffected.
-
-Core design principles:
-
-1. Define backend-agnostic rollout request/response contracts (`RolloutBackendRequest`/`RolloutBackendResponse`).
-2. Isolate protocol differences through backend adapters (`SGLangClient`, `VLLMClient`).
-3. Apply explicit capability gating for non-equivalent features (abort, routed experts, prompt logprobs).
-4. Managed mode -- Slime manages the vLLM process lifecycle within Ray, on par with the SGLang path.
-5. Weight synchronization leverages vLLM's native weight transfer API, automatically selecting the backend based on deployment mode:
-   - **Colocate mode**: CUDA IPC (`IPCWeightTransferEngine`) -- zero-copy via shared GPU memory.
-   - **Non-colocate mode**: NCCL broadcast (`NCCLWeightTransferEngine`) -- direct GPU transfer.
-
-**Current status**: Phase 1 is complete and verified. GRPO training runs successfully with Qwen2.5-0.5B + GSM8K on 4 GPUs in colocate mode.
-
-## 2. Motivation
-
-Slime currently assumes SGLang behavior in multiple places:
-
-- Rollout generation response format (`meta_info.finish_reason.type`, `output_token_logprobs`, optional `routed_experts`)
-- Router control-plane interactions (`/workers`, `/list_workers`, `/abort_request`)
-- Ray-based rollout server startup flow
-
-Supporting vLLM is therefore not just "swapping a URL" -- it requires compatibility layers on both the data plane and control plane.
-
-## 3. Goals and Non-Goals
-
-### Goals
-
-- Add `--rollout-backend {sglang,vllm}`
-- Keep SGLang as the default, unchanged
-- Keep the training-side interface stable
-- Support GRPO rollout with explicit compatibility semantics and observability
-- **Support colocate mode** (training and inference share GPUs, weight sync via CUDA IPC)
-
-### Non-Goals (Current Phase)
-
-- Full feature parity with all SGLang capabilities
-- R3 routing replay on vLLM
-- Multi-instance vLLM + router load balancing
-
-## 4. Architecture and Interface Changes
-
-### 4.1 Architecture Changes
-
-#### End-to-End Architecture Diagram
-
-```text
-                          +--------------------+
-                          |    TrainerLoop     |
-                          +---------+----------+
-                                    |
-              +---------------------+---------------------+
-              | (generate)                    (update_weights)
-              v                                           v
-    +--------------------+              +------------------------------------+
-    |  RolloutFunction   |              | weight updater (auto-selected)     |
-    |  (sglang_rollout)  |              |  colocate → UpdateWeightFromTensor |
-    +---------+----------+              |  otherwise → ...FromDistributed    |
-              |                         +---------------+--------------------+
-              v                                         |
-  +-------------------------------+                     |
-  | RolloutBackendRequest         |                     |
-  +---------------+---------------+                     |
-                  |                                     v
-                  v                       +----------------------------+
-     +--------------------------+         |  VLLMEngine (Ray actor)    |
-     |  RolloutBackendClient    |         |  weight transfer backend:  |
-     +------------+-------------+         |   colocate → IPC           |
-                  |                       |   otherwise → NCCL         |
-    +-------------+-------------+         +-------------+--------------+
-    |                           |                       |
-    v                           v                       v
-+------------+          +-----------+           +---------------+
-| SGLang     |          | VLLM      |           | vLLM server   |
-| Client     |          | Client    |           | /update_weights|
-+-----+------+          +-----+-----+          | /pause /resume |
-      |                       |                 | /sleep /wake_up|
-      v                       v                 +---------------+
- +-----------+        +-------------+
- | SGLang    |        | vLLM server |
- | Router    |        | /v1/compl.  |
- +-----------+        +-------------+
-       \                     /
-        \                   /
-         v                 v
-  +-----------------------------------+
-  | RolloutBackendResponse            |
-  | (text, token_ids, token_logprobs, |
-  |  finish_reason, backend_raw)      |
-  +----------------+------------------+
-                   |
-                   v
-  +-----------------------------------+
-  | SampleUpdate + Training Pipeline  |
-  | (backend-agnostic consumption)    |
-  +-----------------------------------+
-```
-
-#### Component Responsibility Comparison
-
-| Component | Before RFC | After RFC |
-|---|---|---|
-| `RolloutFunction` | Generic rollout + SGLang protocol details | Generic rollout orchestration only |
-| Backend protocol layer | Implicit in rollout logic | Explicit `RolloutBackendClient` adapter |
-| `SGLangClient` | Did not exist (logic scattered) | Owns all SGLang request/response/control-plane details |
-| `VLLMClient` | Did not exist | Owns vLLM `/v1/completions` request/response mapping |
-| Training/sample pipeline | Indirectly consumed SGLang-format fields | Consumes unified contract fields, backend-agnostic |
-| Weight sync | SGLang IPC or NCCL only | Auto-adapts: SGLang IPC / vLLM IPC / vLLM NCCL |
-
-#### Control-Plane Behavior Split
-
-| Control-plane behavior | SGLang path | vLLM path |
-|---|---|---|
-| Worker registration/discovery API | Supported | Not needed |
-| Worker-level abort | Supported (`/abort_request`) | Degraded to timeout/cancel strategy |
-| Routed experts replay | Supported | Explicitly unsupported (capability-gated) |
-| Health check | Existing SGLang/SlimeRouter | `GET /health` direct connection |
-| Memory management (sleep/wake) | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` |
-
-#### A) Rollout Backend Abstraction Layer
-
-- `RolloutBackendClient` (`slime/rollout/backends/base_client.py`): defines `generate()` + `capabilities` abstract interface
-- `BackendCapabilities` dataclass: declares `supports_abort`, `supports_routed_experts`, `supports_prompt_logprobs`
-- `SGLangClient` (extracted from existing code), `VLLMClient` (new)
-
-**Impact**: Rollout logic calls a unified interface and no longer directly embeds SGLang HTTP semantics.
-
-#### B) Rollout Execution Path
-
-- `sglang_rollout.generate` constructs `RolloutBackendRequest` and consumes `RolloutBackendResponse`
-- Selects `SGLangClient` or `VLLMClient` based on `--rollout-backend`
-- `VLLMClient` calls vLLM's OpenAI-compatible `/v1/completions` endpoint
-
-**Impact**: Training-side logic remains stable; backend details are encapsulated in adapters.
-
-#### C) RolloutManager Startup Path
-
-- Existing path (SGLang): unchanged
-- vLLM path: `_start_vllm_rollout_servers()` creates a `VLLMEngine` Ray actor that starts and manages a local vLLM server process
-
-**Impact**: vLLM gets the same lifecycle management as SGLang.
-
-#### D) Weight Synchronization
-
-Weight synchronization supports two modes, automatically selected based on deployment:
-
-**Selection logic** (`actor.py`):
-```python
-update_weight_cls = UpdateWeightFromTensor if args.colocate else UpdateWeightFromDistributed
-```
-Identical to the SGLang selection logic -- backend-agnostic.
-
-##### D.1) Colocate Mode: CUDA IPC
-
-The vLLM server starts with `--weight-transfer-config '{"backend": "ipc"}'`.
-
-Call chain:
-```text
-UpdateWeightFromTensor.update_weights()
-  → engine.pause_generation.remote()           # VLLMEngine → POST /pause?mode=abort
-  → _send_to_colocated_vllm_engine()
-      → each training rank:
-          reduce_tensor(tensor)                 # create CUDA IPC handle
-          {gpu_uuid: ipc_handle}                # keyed by physical GPU UUID
-      → dist.gather_object (Gloo)              # collect handles from all TP ranks
-      → pickle + base64 encode
-      → engine.update_weights_from_tensor.remote()
-          → VLLMEngine → POST /update_weights
-            { "update_info": {
-                "names": [...],
-                "dtype_names": [...],
-                "shapes": [...],
-                "ipc_handles_pickled": "base64..."
-            }}
-          → vLLM IPCWeightTransferEngine.receive_weights()
-            → each TP worker looks up its IPC handle by GPU UUID
-            → func(*args) reconstructs tensor → load_weights()
-  → engine.continue_generation.remote()        # VLLMEngine → POST /resume
-```
-
-Key points:
-- The training process and vLLM workers share the same physical GPU; CUDA IPC enables zero-copy weight transfer
-- For TP>1, IPC handles from all training ranks are merged via Gloo gather; each parameter contains a mapping of all GPU UUIDs
-- The training side must keep tensor references alive until `ray.get()` returns (vLLM has finished reading)
-
-##### D.2) Non-Colocate Mode: NCCL Broadcast
-
-The vLLM server starts with `--weight-transfer-config '{"backend": "nccl"}'`.
-
-Call chain:
-```text
-UpdateWeightFromDistributed.update_weights()
-  → engine.init_weights_update_group.remote()  # VLLMEngine → POST /init_weight_transfer_engine
-      → vLLM NCCLWeightTransferEngine initializes StatelessProcessGroup + PyNcclCommunicator
-  → training side: NCCLWeightTransferEngine.trainer_init()
-      → establishes a matching NCCL communicator with vLLM
-  → PyNcclCommunicator.broadcast(tensor, src=0)  # direct GPU transfer to vLLM workers
-  → engine.update_weights_from_distributed.remote()
-      → VLLMEngine → POST /update_weights
-        { "update_info": { "names": [...], "dtype_names": [...], "shapes": [...] }}
-      → vLLM NCCLWeightTransferEngine.receive_weights()
-```
-
-Key points:
-- The training side uses vLLM's `NCCLWeightTransferEngine.trainer_init()` to initialize NCCL, compatible with vLLM's internal `StatelessProcessGroup` + `PyNcclCommunicator`
-- Cannot use `torch.distributed.init_process_group` (incompatible with vLLM)
-
-#### VLLMEngine Endpoint Mapping
-
-`VLLMEngine` exposes method signatures compatible with `SGLangEngine`. Internal HTTP endpoint mapping:
-
-| Method | VLLMEngine internal HTTP call |
-|---|---|
-| `pause_generation()` | `POST /pause?mode=abort` |
-| `flush_cache()` | no-op |
-| `continue_generation()` | `POST /resume` |
-| `init_weights_update_group(...)` | `POST /init_weight_transfer_engine` |
-| `update_weights_from_distributed(...)` | `POST /update_weights` (NCCL mode) |
-| `update_weights_from_tensor(...)` | `POST /update_weights` (IPC mode) |
-| `release_memory_occupation()` | `POST /sleep?level=1&mode=abort` |
-| `resume_memory_occupation()` | `POST /wake_up` |
-| `health_generate()` | `GET /health` |
-| `shutdown()` | `process.terminate()` |
-
-**Impact**: Training-side code (`UpdateWeightFromTensor`, `UpdateWeightFromDistributed`, actor code) uses a unified interface, unaware of the specific backend.
-
-#### E) Router Decision
-
-vLLM does not use a router in the current phase:
-- Single vLLM instance; `VLLMClient` connects directly to the local vLLM server port
-- SGLang Model Gateway only supports SGLang workers, not applicable
-- SlimeRouter is only needed for R3 / radix-tree caching scenarios
-
-### 4.2 Interface Changes
-
-#### A) New CLI Arguments
-
-- `--rollout-backend {sglang,vllm}` -- select rollout backend
-- `--vllm-base-url` -- manually specify vLLM server address (not needed when auto-managed)
-- `--vllm-model` -- model path for vLLM to load (defaults to `--hf-checkpoint`)
-- `--vllm-max-retries` -- max retries for generation requests
-- `--vllm-enforce-eager` -- disable CUDA graph (default True)
-
-#### B) New Internal Unified Interface
-
-Added to `slime/rollout/base_types.py`:
-
-- `RolloutBackendRequest`: input_ids, sampling_params, return_logprob, return_routed_experts, image_data, session_id
-- `RolloutBackendResponse`: text, output_token_ids, output_token_logprobs, finish_reason (`stop|length|abort`), prompt_tokens, completion_tokens, backend_raw, routed_experts
-
-#### C) Capability Gating
-
-`BackendCapabilities` dataclass declares capabilities per backend:
-
-| Capability | SGLangClient | VLLMClient |
-|---|---|---|
-| `supports_abort` | True | False |
-| `supports_routed_experts` | True | False |
-| `supports_prompt_logprobs` | True | False |
-
-Unsupported features are explicitly gated, logged, or fail fast at call time.
-
-## 5. File-Level Change List
-
-### New Files
-
-| File | Description |
-|---|---|
-| `slime/rollout/backends/base_client.py` | `RolloutBackendClient` abstract interface + `BackendCapabilities` |
-| `slime/rollout/backends/sglang_client.py` | SGLang rollout client (extracted from existing code) |
-| `slime/rollout/backends/vllm_client.py` | vLLM rollout client (`/v1/completions` adapter) |
-| `slime/rollout/backends/__init__.py` | Adapter exports |
-| `slime/backends/vllm_utils/vllm_engine.py` | `VLLMEngine` Ray actor (process management + weight sync) |
-| `run-qwen2.5-0.5B-vllm.sh` | vLLM validation script (Qwen2.5-0.5B + GSM8K + colocate) |
-
-### Modified Files
-
-| File | Description |
-|---|---|
-| `slime/rollout/base_types.py` | Added `RolloutBackendRequest` / `RolloutBackendResponse` |
-| `slime/rollout/sglang_rollout.py` | Uses backend adapter instead of hard-coded SGLang protocol |
-| `slime/utils/arguments.py` | Added `--rollout-backend`, vLLM arguments; sets sglang alias defaults for vLLM |
-| `slime/ray/rollout.py` | `_start_vllm_rollout_servers()` to start VLLMEngine |
-| `slime/backends/megatron_utils/actor.py` | Weight updater selection: `colocate → UpdateWeightFromTensor` (backend-agnostic) |
-| `slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py` | Added `_send_to_colocated_vllm_engine()` (CUDA IPC path) |
-| `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` | vLLM NCCL compatibility: `NCCLWeightTransferEngine.trainer_init()` + `PyNcclCommunicator.broadcast()` |
-
-### Unchanged Files
-
-| File | Description |
-|---|---|
-| `train.py` | Training loop is backend-agnostic |
-| SGLang-related files | SGLang path is completely unaffected |
-
-### Not Used (Differences from Initial RFC Draft)
-
-- **No `worker_extension.py`**: vLLM's native weight transfer API fully meets requirements; no custom WorkerExtension needed
-- **No `collective_rpc`**: Uses vLLM's `/init_weight_transfer_engine` and `/update_weights` endpoints
-
-## 6. Capability Compatibility Matrix
-
-| Capability | SGLang | vLLM | Handling |
-|---|---|---|---|
-| Token-level response logprobs | Supported | Supported (`choice.logprobs.token_logprobs`) | Adapter normalization |
-| Output token IDs | `meta_info` | `choice.token_ids` | Adapter normalization |
-| Finish reason | `stop\|length\|abort` | `stop\|length` + others | Canonical mapping |
-| Worker-level abort | Supported | Not supported | Timeout/cancel degradation |
-| Prompt logprobs (OPD) | Supported | Partial | Capability-gated |
-| Routed experts (R3) | Supported | Not supported | Explicit gate |
-| Colocate (GPU sharing) | Supported (FlattenedTensorBucket IPC) | Supported (CUDA IPC handles) | Native IPC per backend |
-| Memory management | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | Unified interface |
-| `include_stop_str_in_output` | `no_stop_trim` | `include_stop_str_in_output` | Parameter mapping |
-
-## 7. Phased Plan
-
-### Phase 1: Managed vLLM GRPO Training ✅ Done
-
-**Goal**: Qwen2.5-0.5B GRPO training on 4 GPUs in colocate mode with GSM8K dataset, running successfully.
-
-Completed deliverables:
-- `VLLMEngine` Ray actor: process lifecycle management, `/pause`/`/resume`/`/sleep`/`/wake_up`
-- `VLLMClient` rollout adapter: `/v1/completions` request mapping, logprob/finish_reason/token_ids normalization
-- `RolloutBackendRequest`/`Response` unified contract
-- `SGLangClient` extraction
-- `_start_vllm_rollout_servers` startup path
-- Argument parsing branching (`--rollout-backend vllm`, sglang alias defaults)
-- **Colocate mode**: CUDA IPC weight transfer (`_send_to_colocated_vllm_engine` → `IPCWeightTransferEngine`)
-- **Non-colocate mode**: NCCL weight transfer (`NCCLWeightTransferEngine.trainer_init()` → `PyNcclCommunicator.broadcast()`)
-- Validation script: `run-qwen2.5-0.5B-vllm.sh`
-
-Technical decisions:
-- No router (single vLLM instance, direct connection)
-- Colocate weight sync: CUDA IPC (training and vLLM share GPU memory, zero-copy)
-- Non-colocate weight sync: vLLM native NCCL (`StatelessProcessGroup` + `PyNcclCommunicator`)
-- Rollout function: reuses `sglang_rollout.py` (backend-agnostic orchestration)
-- Training-side weight updater selection is identical to SGLang (`colocate → UpdateWeightFromTensor`)
-
-Acceptance results:
-- ✅ Qwen2.5-0.5B + GSM8K + 4 GPU colocate training runs successfully
-- ✅ Weight sync correctness: rollout outputs change with each training update
-- ✅ Existing SGLang path unaffected
-
-### Phase 2: Multi-Instance vLLM + Async Training
-
-**Goal**: Support multi-instance vLLM rollout with [vllm-project/router](https://github.com/vllm-project/router) for load balancing, and verify async training (`train_async.py`) on larger models.
-
-Key deliverables:
-- Integrate [vllm-project/router](https://github.com/vllm-project/router) as the vLLM-side load balancer
-- `start_rollout_servers` launches N `VLLMEngine` actors + one vllm-router process
-- `VLLMClient` generation requests routed through vllm-router instead of direct connection
-- Verify async training (`train_async.py`) correctness with the vLLM backend
-
-Acceptance criteria:
-- Multi-instance (e.g., 2-4 vLLM workers) rollout stable over 20+ rollout steps
-- Async training (`train_async.py`) with no hangs or weight sync race conditions
-- Throughput improvement compared to single-instance baseline
-- Larger model verification (e.g., Qwen3-4B, TP=2)
-
-### Phase 3 (Future): Advanced Features
-
-- Abort/cancel strategy refinement
-- Prompt logprob support (OPD scenarios)
-- Deterministic computation verification
-- Performance benchmarks (vLLM vs SGLang throughput comparison)
-
-## 8. Validation Strategy
-
-### Completed (Phase 1)
-
-- ✅ End-to-end GRPO training: `run-qwen2.5-0.5B-vllm.sh` (Qwen2.5-0.5B, GSM8K, 4 GPUs, colocate)
-- ✅ Weight sync correctness verification
-- ✅ SGLang path non-regression
-
-### Remaining
-
-- Unit tests: finish reason normalization, token/logprob alignment, capability-gated behavior
-- Integration tests: SGLang vs vLLM response schema comparison
-- Stress tests: higher concurrency / latency pressure
-
-## 9. Key Implementation Challenges and Solutions
-
-### 9.1 vLLM NCCL Incompatibility
-
-**Problem**: vLLM internally uses `StatelessProcessGroup` + `PyNcclCommunicator` for NCCL communication, which is incompatible with `torch.distributed.init_process_group`.
-
-**Solution**: The training side uses `NCCLWeightTransferEngine.trainer_init()` to initialize NCCL, ensuring a matching communicator is established with the vLLM side.
-
-### 9.2 vLLM Logprobs Format Differences
-
-**Problem**: vLLM's `/v1/completions` logprobs format (`choice.token_ids`, `choice.logprobs.token_logprobs`) differs from SGLang's.
-
-**Solution**: `VLLMClient` performs explicit mapping, converting vLLM's format into the unified `RolloutBackendResponse` format.
diff --git a/rfc-vllm-rollout-backend.md b/rfc-vllm-rollout-backend.md
deleted file mode 100644
index fc8e569793..0000000000
--- a/rfc-vllm-rollout-backend.md
+++ /dev/null
@@ -1,436 +0,0 @@
-# RFC: 在 Slime 中支持 vLLM 作为 Rollout Backend
-
-## 1. 概要
-
-本 RFC 提议在 Slime 中增加 **vLLM** 作为一等 rollout backend，同时保持现有 SGLang 行为不变，不影响 GRPO 训练流程。
-
-核心设计思路：
-
-1. 定义 backend 无关的 rollout 请求/响应契约（`RolloutBackendRequest`/`RolloutBackendResponse`）。
-2. 通过 backend adapter（`SGLangClient`、`VLLMClient`）隔离协议差异。
-3. 对不等价能力（abort、routed experts、prompt logprobs）做显式 capability gating。
-4. Managed 模式 —— Slime 在 Ray 内管理 vLLM 进程生命周期，与 SGLang 路径同等级别。
-5. 权重同步利用 vLLM 原生 weight transfer API，根据部署模式自动选择后端：
-   - **Colocate 模式**：CUDA IPC（`IPCWeightTransferEngine`），GPU 共享内存零拷贝。
-   - **Non-colocate 模式**：NCCL broadcast（`NCCLWeightTransferEngine`），GPU 直传。
-
-**当前状态**：Phase 1 已完成并验证通过。Qwen2.5-0.5B + GSM8K + 4 GPU colocate 模式下 GRPO 训练正常运行。
-
-## 2. 为什么需要这个
-
-Slime 当前在多处假设 SGLang 行为：
-
-- rollout 生成响应格式（`meta_info.finish_reason.type`、`output_token_logprobs`、可选 `routed_experts`）
-- router 控制面交互（`/workers`、`/list_workers`、`/abort_request`）
-- Ray 内 rollout server 启动流程
-
-因此支持 vLLM 不是"换个 URL"，而是需要在数据面和控制面做兼容层。
-
-## 3. 目标与非目标
-
-### 目标
-
-- 新增 `--rollout-backend {sglang,vllm}`
-- SGLang 路径默认不变
-- 训练侧接口保持稳定
-- 支持 GRPO rollout，具有显式兼容性语义和可观测性
-- **支持 colocate 模式**（训练与推理共享 GPU，通过 CUDA IPC 同步权重）
-
-### 非目标（当前阶段）
-
-- 所有 SGLang 特性的完全对等
-- vLLM 上的 R3 路由回放
-- 多 vLLM 实例 + router 负载均衡
-
-## 4. 架构和接口改动
-
-### 4.1 架构改动
-
-#### 端到端架构图
-
-```text
-                          +--------------------+
-                          |    TrainerLoop     |
-                          +---------+----------+
-                                    |
-              +---------------------+---------------------+
-              | (generate)                    (update_weights)
-              v                                           v
-    +--------------------+              +------------------------------------+
-    |  RolloutFunction   |              | weight updater (自动选择)            |
-    |  (sglang_rollout)  |              |  colocate → UpdateWeightFromTensor  |
-    +---------+----------+              |  otherwise → ...FromDistributed     |
-              |                         +---------------+--------------------+
-              v                                         |
-  +-------------------------------+                     |
-  | RolloutBackendRequest         |                     |
-  +---------------+---------------+                     |
-                  |                                     v
-                  v                       +----------------------------+
-     +--------------------------+         |  VLLMEngine (Ray actor)    |
-     |  RolloutBackendClient    |         |  weight transfer backend:  |
-     +------------+-------------+         |   colocate → IPC           |
-                  |                       |   otherwise → NCCL         |
-    +-------------+-------------+         +-------------+--------------+
-    |                           |                       |
-    v                           v                       v
-+------------+          +-----------+           +---------------+
-| SGLang     |          | VLLM      |           | vLLM server   |
-| Client     |          | Client    |           | /update_weights|
-+-----+------+          +-----+-----+          | /pause /resume |
-      |                       |                 | /sleep /wake_up|
-      v                       v                 +---------------+
- +-----------+        +-------------+
- | SGLang    |        | vLLM server |
- | Router    |        | /v1/compl.  |
- +-----------+        +-------------+
-       \                     /
-        \                   /
-         v                 v
-  +-----------------------------------+
-  | RolloutBackendResponse            |
-  | (text, token_ids, token_logprobs, |
-  |  finish_reason, backend_raw)      |
-  +----------------+------------------+
-                   |
-                   v
-  +-----------------------------------+
-  | SampleUpdate + Training Pipeline  |
-  | (backend 无关消费)                 |
-  +-----------------------------------+
-```
-
-#### 组件职责对照
-
-| 组件 | RFC 前 | RFC 后 |
-|---|---|---|
-| `RolloutFunction` | 包含通用 rollout + SGLang 协议细节 | 只包含通用 rollout 编排 |
-| Backend 协议层 | 隐含在 rollout 逻辑里 | 显式的 `RolloutBackendClient` adapter |
-| `SGLangClient` | 不存在（逻辑分散） | 拥有 SGLang 请求/响应/控制面的全部细节 |
-| `VLLMClient` | 不存在 | 拥有 vLLM `/v1/completions` 请求/响应映射 |
-| 训练/sample 管线 | 间接消费 SGLang 格式字段 | 消费统一契约字段，backend 无关 |
-| 权重同步 | 仅 SGLang IPC 或 NCCL | 自动适配：SGLang IPC / vLLM IPC / vLLM NCCL |
-
-#### 控制面行为拆分
-
-| 控制面行为 | SGLang 路径 | vLLM 路径 |
-|---|---|---|
-| Worker 注册/发现 API | 支持 | 不需要 |
-| Worker 级 abort | 支持（`/abort_request`） | 降级为超时/取消策略 |
-| Routed experts 回放 | 支持 | 显式不支持（capability-gated） |
-| 健康检查 | 现有 SGLang/SlimeRouter | `GET /health` 直连 |
-| 内存管理 (sleep/wake) | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` |
-
-#### A) Rollout backend 抽象层
-
-- `RolloutBackendClient`（`slime/rollout/backends/base_client.py`）：定义 `generate()` + `capabilities` 抽象接口
-- `BackendCapabilities` dataclass：声明 `supports_abort`、`supports_routed_experts`、`supports_prompt_logprobs`
-- `SGLangClient`（从现有代码提取）、`VLLMClient`（新建）
-
-**影响**：rollout 逻辑调用统一接口，不再直接嵌入 SGLang HTTP 语义。
-
-#### B) Rollout 执行路径
-
-- `sglang_rollout.generate` 构造 `RolloutBackendRequest`，消费 `RolloutBackendResponse`
-- 根据 `--rollout-backend` 选择 `SGLangClient` 或 `VLLMClient`
-- `VLLMClient` 调用 vLLM OpenAI-compatible `/v1/completions` 端点
-
-**影响**：训练侧逻辑保持稳定，backend 细节移入 adapter。
-
-#### C) RolloutManager 启动路径
-
-- 现有路径（SGLang）：保持不变
-- vLLM 路径：`_start_vllm_rollout_servers()` 创建 `VLLMEngine` Ray actor，启动管理本地 vLLM server 进程
-
-**影响**：vLLM 获得与 SGLang 相同的生命周期管理。
-
-#### D) 权重同步
-
-权重同步支持两种模式，根据部署方式自动选择：
-
-**模式选择逻辑**（`actor.py`）：
-```python
-update_weight_cls = UpdateWeightFromTensor if args.colocate else UpdateWeightFromDistributed
-```
-与 SGLang 使用完全相同的选择逻辑，backend 无关。
-
-##### D.1) Colocate 模式：CUDA IPC
-
-vLLM server 启动时配置 `--weight-transfer-config '{"backend": "ipc"}'`。
-
-调用链：
-```text
-UpdateWeightFromTensor.update_weights()
-  → engine.pause_generation.remote()           # VLLMEngine → POST /pause?mode=abort
-  → _send_to_colocated_vllm_engine()
-      → 每个训练 rank:
-          reduce_tensor(tensor)                 # 创建 CUDA IPC handle
-          {gpu_uuid: ipc_handle}                # 以物理 GPU UUID 为 key
-      → dist.gather_object (Gloo)              # 收集所有 TP rank 的 handles
-      → pickle + base64 编码
-      → engine.update_weights_from_tensor.remote()
-          → VLLMEngine → POST /update_weights
-            { "update_info": {
-                "names": [...],
-                "dtype_names": [...],
-                "shapes": [...],
-                "ipc_handles_pickled": "base64..."
-            }}
-          → vLLM IPCWeightTransferEngine.receive_weights()
-            → 每个 TP worker 根据自己的 GPU UUID 查找 IPC handle
-            → func(*args) 重建 tensor → load_weights()
-  → engine.continue_generation.remote()        # VLLMEngine → POST /resume
-```
-
-关键点：
-- 训练进程和 vLLM worker 共享同一物理 GPU，CUDA IPC 实现零拷贝权重传输
-- TP>1 时，各 training rank 的 IPC handles 通过 Gloo gather 合并，每个参数包含所有 GPU UUID 的映射
-- 训练侧必须保持 tensor 引用直到 `ray.get()` 返回（vLLM 读取完毕）
-
-##### D.2) Non-colocate 模式：NCCL broadcast
-
-vLLM server 启动时配置 `--weight-transfer-config '{"backend": "nccl"}'`。
-
-调用链：
-```text
-UpdateWeightFromDistributed.update_weights()
-  → engine.init_weights_update_group.remote()  # VLLMEngine → POST /init_weight_transfer_engine
-      → vLLM NCCLWeightTransferEngine 初始化 StatelessProcessGroup + PyNcclCommunicator
-  → 训练侧: NCCLWeightTransferEngine.trainer_init()
-      → 与 vLLM 建立匹配的 NCCL 通信组
-  → PyNcclCommunicator.broadcast(tensor, src=0)  # GPU 直传到 vLLM worker
-  → engine.update_weights_from_distributed.remote()
-      → VLLMEngine → POST /update_weights
-        { "update_info": { "names": [...], "dtype_names": [...], "shapes": [...] }}
-      → vLLM NCCLWeightTransferEngine.receive_weights()
-```
-
-关键点：
-- 训练侧使用 vLLM 的 `NCCLWeightTransferEngine.trainer_init()` 初始化 NCCL，与 vLLM 内部的 `StatelessProcessGroup` + `PyNcclCommunicator` 兼容
-- 不能使用 `torch.distributed.init_process_group`（vLLM 不兼容）
-
-#### VLLMEngine 端点映射
-
-`VLLMEngine` 暴露与 `SGLangEngine` 兼容的方法签名。内部 HTTP 端点映射：
-
-| 方法 | VLLMEngine 内部 HTTP 调用 |
-|---|---|
-| `pause_generation()` | `POST /pause?mode=abort` |
-| `flush_cache()` | no-op |
-| `continue_generation()` | `POST /resume` |
-| `init_weights_update_group(...)` | `POST /init_weight_transfer_engine` |
-| `update_weights_from_distributed(...)` | `POST /update_weights` (NCCL 模式) |
-| `update_weights_from_tensor(...)` | `POST /update_weights` (IPC 模式) |
-| `release_memory_occupation()` | `POST /sleep?level=1&mode=abort` |
-| `resume_memory_occupation()` | `POST /wake_up` |
-| `health_generate()` | `GET /health` |
-| `shutdown()` | `process.terminate()` |
-
-**影响**：训练侧代码（`UpdateWeightFromTensor`、`UpdateWeightFromDistributed`、actor 代码）使用统一接口，不感知具体 backend。
-
-#### E) Router 决策
-
-vLLM 当前阶段不使用 router：
-- 单 vLLM 实例，`VLLMClient` 直连本地 vLLM server 端口
-- SGLang Model Gateway 只支持 SGLang worker，不适用
-- SlimeRouter 仅在 R3 / radix-tree caching 时需要
-
-### 4.2 接口改动
-
-#### A) 新增 CLI 接口
-
-- `--rollout-backend {sglang,vllm}` —— 选择 rollout backend
-- `--vllm-base-url` —— 手动指定 vLLM server 地址（自动管理时无需设置）
-- `--vllm-model` —— vLLM 加载的模型路径（默认同 `--hf-checkpoint`）
-- `--vllm-max-retries` —— 生成请求最大重试次数
-- `--vllm-enforce-eager` —— 是否禁用 CUDA graph（默认 True）
-
-#### B) 新增内部统一接口
-
-`slime/rollout/base_types.py` 中新增：
-
-- `RolloutBackendRequest`：input_ids、sampling_params、return_logprob、return_routed_experts、image_data、session_id
-- `RolloutBackendResponse`：text、output_token_ids、output_token_logprobs、finish_reason（`stop|length|abort`）、prompt_tokens、completion_tokens、backend_raw、routed_experts
-
-#### C) 能力门控
-
-`BackendCapabilities` dataclass 声明各 backend 支持的能力：
-
-| 能力 | SGLangClient | VLLMClient |
-|---|---|---|
-| `supports_abort` | True | False |
-| `supports_routed_experts` | True | False |
-| `supports_prompt_logprobs` | True | False |
-
-不支持的特性在调用时显式 gate、记录日志、或 fail fast。
-
-## 5. 文件级改动清单
-
-### 新增文件
-
-| 文件 | 说明 |
-|---|---|
-| `slime/rollout/backends/base_client.py` | `RolloutBackendClient` 抽象接口 + `BackendCapabilities` |
-| `slime/rollout/backends/sglang_client.py` | SGLang rollout client（从现有代码提取） |
-| `slime/rollout/backends/vllm_client.py` | vLLM rollout client（`/v1/completions` 适配） |
-| `slime/rollout/backends/__init__.py` | adapter 导出 |
-| `slime/backends/vllm_utils/vllm_engine.py` | `VLLMEngine` Ray actor（进程管理 + 权重同步） |
-| `run-qwen2.5-0.5B-vllm.sh` | vLLM 验证脚本（Qwen2.5-0.5B + GSM8K + colocate） |
-
-### 修改文件
-
-| 文件 | 说明 |
-|---|---|
-| `slime/rollout/base_types.py` | 新增 `RolloutBackendRequest` / `RolloutBackendResponse` |
-| `slime/rollout/sglang_rollout.py` | 使用 backend adapter 代替硬编码 SGLang 协议 |
-| `slime/utils/arguments.py` | 新增 `--rollout-backend`、vLLM 参数；vLLM 时设置 sglang 别名默认值 |
-| `slime/ray/rollout.py` | `_start_vllm_rollout_servers()` 启动 VLLMEngine |
-| `slime/backends/megatron_utils/actor.py` | 权重更新器选择：`colocate → UpdateWeightFromTensor`（backend 无关） |
-| `slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py` | 新增 `_send_to_colocated_vllm_engine()`（CUDA IPC 路径） |
-| `slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py` | vLLM NCCL 兼容：`NCCLWeightTransferEngine.trainer_init()` + `PyNcclCommunicator.broadcast()` |
-
-### 不变文件
-
-| 文件 | 说明 |
-|---|---|
-| `train.py` | 训练循环不感知 backend |
-| SGLang 相关文件 | SGLang 路径完全不受影响 |
-
-### 未使用（与 RFC 初稿的差异）
-
-- **无 `worker_extension.py`**：vLLM 原生 weight transfer API 完全满足需求，无需自定义 WorkerExtension
-- **无 `collective_rpc`**：使用 vLLM 的 `/init_weight_transfer_engine` 和 `/update_weights` 端点
-
-## 6. 能力兼容矩阵
-
-| 能力 | SGLang | vLLM | 处理方式 |
-|---|---|---|---|
-| Token 级响应 logprobs | 支持 | 支持（`choice.logprobs.token_logprobs`） | adapter 归一化 |
-| Output token IDs | `meta_info` | `choice.token_ids` | adapter 归一化 |
-| Finish reason | `stop\|length\|abort` | `stop\|length` + 其他 | canonical 映射 |
-| Worker 级 abort | 支持 | 不支持 | 超时/取消降级 |
-| Prompt logprobs（OPD） | 支持 | 部分 | capability-gated |
-| Routed experts（R3） | 支持 | 不支持 | 显式 gate |
-| Colocate（GPU 共享） | 支持（FlattenedTensorBucket IPC） | 支持（CUDA IPC handles） | 各自原生 IPC |
-| 内存管理 | `release/resume_memory_occupation` | `POST /sleep` / `POST /wake_up` | 接口统一 |
-| `include_stop_str_in_output` | `no_stop_trim` | `include_stop_str_in_output` | 参数映射 |
-
-## 7. 分阶段计划
-
-### 第一阶段：Managed vLLM GRPO 训练 ✅ Done
-
-**目标**：Qwen2.5-0.5B GRPO 4 卡 colocate 训练，GSM8K 数据集，训练正常运行。
-
-已完成交付物：
-- `VLLMEngine` Ray actor：进程生命周期管理、`/pause`/`/resume`/`/sleep`/`/wake_up`
-- `VLLMClient` rollout adapter：`/v1/completions` 请求映射、logprob/finish_reason/token_ids 归一化
-- `RolloutBackendRequest`/`Response` 统一契约
-- `SGLangClient` 提取
-- `_start_vllm_rollout_servers` 启动路径
-- 参数解析分流（`--rollout-backend vllm`，sglang 别名默认值）
-- **Colocate 模式**：CUDA IPC 权重传输（`_send_to_colocated_vllm_engine` → `IPCWeightTransferEngine`）
-- **Non-colocate 模式**：NCCL 权重传输（`NCCLWeightTransferEngine.trainer_init()` → `PyNcclCommunicator.broadcast()`）
-- 验证脚本：`run-qwen2.5-0.5B-vllm.sh`
-
-技术决策：
-- 不使用 router（单 vLLM 实例，直连）
-- Colocate 权重同步：CUDA IPC（训练和 vLLM 共享 GPU 内存，零拷贝）
-- Non-colocate 权重同步：vLLM 原生 NCCL（`StatelessProcessGroup` + `PyNcclCommunicator`）
-- Rollout 函数：复用 `sglang_rollout.py`（backend 无关编排）
-- 训练侧权重更新器选择与 SGLang 一致（`colocate → UpdateWeightFromTensor`）
-
-验收结果：
-- ✅ Qwen2.5-0.5B + GSM8K + 4 GPU colocate 训练运行通过
-- ✅ 权重同步正确性：每轮 rollout 输出随训练更新变化
-- ✅ 现有 SGLang 路径不受影响
-
-### 第二阶段：多 vLLM 实例 + 异步训推
-
-**目标**：支持多 vLLM 实例 rollout，使用 [vllm-project/router](https://github.com/vllm-project/router) 做负载均衡，以异步训推（`train_async.py`）在更大模型上验证。
-
-关键交付物：
-- 集成 [vllm-project/router](https://github.com/vllm-project/router) 作为 vLLM 侧负载均衡器
-- `start_rollout_servers` 拉起 N 个 `VLLMEngine` actor + 一个 vllm-router 进程
-- `VLLMClient` 生成请求从直连改为走 vllm-router
-- 验证异步训推（`train_async.py`）在 vLLM backend 下的正确性
-
-验收标准：
-- 多实例（如 2-4 个 vLLM worker）rollout 在 20+ rollout step 下稳定
-- 异步训推（`train_async.py`）无 hang、无权重同步竞态
-- 相比单实例基线有吞吐提升
-- 更大模型验证（如 Qwen3-4B，TP=2）
-
-### 第三阶段（未来）：高级特性
-
-- Abort/cancel 策略完善
-- Prompt logprob 支持（OPD 场景）
-- 确定性计算验证
-- 性能基准测试（vLLM vs SGLang 吞吐量对比）
-
-## 8. 验证策略
-
-### 已完成（Phase 1）
-
-- ✅ 端到端 GRPO 训练：`run-qwen2.5-0.5B-vllm.sh`（Qwen2.5-0.5B, GSM8K, 4 GPU, colocate）
-- ✅ 权重同步正确性验证
-- ✅ SGLang 路径非回归
-
-### 待完成
-
-- 单元测试：finish reason 归一化、token/logprob 对齐、capability-gated 行为
-- 集成测试：SGLang vs vLLM 响应 schema 对比
-- 压测：更高并发 / 延迟压力
-
-## 9. 实现中遇到的关键问题及解决
-
-### 9.1 vLLM NCCL 不兼容
-
-**问题**：vLLM 内部使用 `StatelessProcessGroup` + `PyNcclCommunicator` 管理 NCCL 通信，与 `torch.distributed.init_process_group` 不兼容。
-
-**解决**：训练侧使用 `NCCLWeightTransferEngine.trainer_init()` 初始化 NCCL，确保与 vLLM 端建立匹配的通信组。
-
-### 9.2 vLLM logprobs 格式差异
-
-**问题**：vLLM `/v1/completions` 的 logprobs 格式（`choice.token_ids`、`choice.logprobs.token_logprobs`）与 SGLang 不同。
-
-**解决**：`VLLMClient` 中做显式映射，将 vLLM 格式转换为 `RolloutBackendResponse` 统一格式。对于缺失 `token_ids` 的情况，回退到 tokenizer。
-
-### 9.3 Colocate 模式下的 IPC 格式差异
-
-**问题**：SGLang 使用 `FlattenedTensorBucket` + `MultiprocessingSerializer` 做 IPC，vLLM 使用独立的 CUDA IPC handles（`reduce_tensor`）。两者格式不兼容。
-
-**解决**：新增 `_send_to_colocated_vllm_engine()` 函数，直接使用 `torch.multiprocessing.reductions.reduce_tensor` 创建 CUDA IPC handles，以 GPU UUID 为 key。TP>1 时通过 Gloo gather 合并各 rank 的 handles，然后 pickle + base64 编码后通过 `VLLMEngine` 转发到 vLLM 的 `/update_weights` 端点。
-
-### 9.4 缺失的 sglang 参数别名
-
-**问题**：vLLM backend 跳过 `sglang_validate_args()`，导致 `sglang_dp_size` 等别名未设置，后续代码 `AttributeError`。
-
-**解决**：在 `arguments.py` 中为 vLLM backend 添加条件分支，设置 `sglang_dp_size`、`sglang_pp_size`、`sglang_ep_size`、`sglang_tp_size` 等默认值。
-
-## 10. 风险与缓解
-
-- **Logprob 语义不匹配**：严格 adapter 检查 + 归一化 + 测试
-- **Abort 不匹配**：能力模型 + 超时/取消降级 + 显式日志
-- **训练漂移**：强制 A/B 对照运行，固定 seed 和配置
-- **IPC handle 生命周期**：训练侧通过 `kept_alive` 列表 + `ray.get()` 阻塞确保 tensor 在 vLLM 读取完成前不被回收
-
-## 11. 向 Slime 维护者提出的开放问题
-
-1. 是否需要在 Phase 2 中将 vLLM 多实例 + router 作为默认部署模式？
-2. 声称支持 GRPO 所需的最低 logprob 保真度是什么？
-3. 是否应在完全对等前阻止 vLLM 的 OPD prompt-logprob 路径？
-4. 推荐为默认 backend 之前需要满足什么质量门槛？
-5. vLLM colocate IPC 路径是否需要支持混合模式（部分 colocate + 部分 distributed）？
-
-## 12. 请求决策
-
-Phase 1 已实现并验证：
-
-- ✅ Contract-first adapter 架构
-- ✅ Managed vLLM 模式（与 SGLang 同等生命周期管理）
-- ✅ 双模式权重同步（colocate: CUDA IPC / non-colocate: NCCL）
-- ✅ 显式 capability-gated 行为
-- ✅ 严格 SGLang 非回归
-
-请批准进入 Phase 2：多 vLLM 实例 + router + 异步训推。
-
diff --git a/run-qwen2.5-0.5B-reproducibility-noncolocate.sh b/run-qwen2.5-0.5B-reproducibility-noncolocate.sh
deleted file mode 100644
index 5499949340..0000000000
--- a/run-qwen2.5-0.5B-reproducibility-noncolocate.sh
+++ /dev/null
@@ -1,137 +0,0 @@
-#!/bin/bash
-# Non-colocate version of run-qwen2.5-0.5B-reproducibility.sh
-# 2 GPUs: 1 for training, 1 for SGLang rollout
-
-# for rerun the task
-pkill -9 sglang
-sleep 3
-ray stop --force
-pkill -9 ray
-pkill -9 python
-sleep 3
-pkill -9 ray
-pkill -9 python
-
-set -ex
-
-export PYTHONBUFFERED=16
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
-source "${SCRIPT_DIR}/scripts/models/qwen2.5-0.5B.sh"
-
-CKPT_ARGS=(
-   --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/
-   --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/
-)
-
-ROLLOUT_ARGS=(
-   --prompt-data /root/gsm8k/train.parquet
-   --input-key messages
-   --label-key label
-   --apply-chat-template
-   --rollout-shuffle
-   --rm-type math
-   --num-rollout 100
-   --rollout-batch-size 32
-   --n-samples-per-prompt 8
-   --rollout-max-response-len 1024
-   --rollout-temperature 1
-
-   --global-batch-size 256
-)
-
-EVAL_ARGS=(
-   --eval-interval 20
-   --eval-prompt-data gsm8k /root/gsm8k/test.parquet
-   --n-samples-per-eval-prompt 1
-   --eval-max-response-len 1024
-   --eval-top-k 1
-)
-
-PERF_ARGS=(
-   --tensor-model-parallel-size 1
-   --sequence-parallel
-   --pipeline-model-parallel-size 1
-   --context-parallel-size 1
-   --expert-model-parallel-size 1
-   --expert-tensor-parallel-size 1
-
-   --use-dynamic-batch-size
-   --max-tokens-per-gpu 9216
-)
-
-GRPO_ARGS=(
-   --advantage-estimator grpo
-   --use-kl-loss
-   --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
-)
-
-OPTIMIZER_ARGS=(
-   --optimizer adam
-   --lr 1e-6
-   --lr-decay-style constant
-   --weight-decay 0.1
-   --adam-beta1 0.9
-   --adam-beta2 0.98
-)
-
-WANDB_ARGS=(
-   --use-wandb
-   --wandb-host https://wandb.ai/
-   --wandb-entity samithuang
-   --wandb-project slime-rl
-   --wandb-group qwen2.5-0.5B-gsm8k-noncolocate
-)
-
-SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 1
-   --sglang-mem-fraction-static 0.7
-
-   --sglang-enable-deterministic-inference
-   --sglang-attention-backend flashinfer
-
-   --deterministic-mode
-)
-
-MISC_ARGS=(
-   --attention-dropout 0.0
-   --hidden-dropout 0.0
-   --accumulate-allreduce-grads-in-fp32
-   --attention-softmax-in-fp32
-   --attention-backend flash
-)
-
-ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats
-
-ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json='{
-     "env_vars": {
-        "PYTHONPATH": "/root/Megatron-LM",
-        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-        "NCCL_ALGO": "Ring",
-        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
-        "CUBLAS_WORKSPACE_CONFIG": ":4096:8"
-     }
-   }' \
-   -- python3 train.py \
-   --actor-num-nodes 1 \
-   --actor-num-gpus-per-node 1 \
-   --num-gpus-per-node 2 \
-   --rollout-num-gpus 1 \
-   --calculate-per-token-loss \
-   --use-slime-router \
-   ${MODEL_ARGS[@]} \
-   ${CKPT_ARGS[@]} \
-   ${ROLLOUT_ARGS[@]} \
-   ${OPTIMIZER_ARGS[@]} \
-   ${GRPO_ARGS[@]} \
-   ${WANDB_ARGS[@]} \
-   ${PERF_ARGS[@]} \
-   ${EVAL_ARGS[@]} \
-   ${SGLANG_ARGS[@]} \
-   ${MISC_ARGS[@]}
diff --git a/run-qwen2.5-0.5B-vllm.sh b/run-qwen2.5-0.5B-vllm.sh
deleted file mode 100644
index 397fac1c58..0000000000
--- a/run-qwen2.5-0.5B-vllm.sh
+++ /dev/null
@@ -1,146 +0,0 @@
-#!/bin/bash
-# vLLM rollout backend validation script (Phase 1)
-# Based on run-qwen2.5-0.5B-reproducibility.sh
-
-# for rerun the task
-pkill -9 vllm
-pkill -9 sglang
-sleep 3
-ray stop --force
-pkill -9 ray
-pkill -9 python
-sleep 3
-pkill -9 ray
-pkill -9 python
-
-
-set -ex
-
-export PYTHONBUFFERED=16
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
-source "${SCRIPT_DIR}/scripts/models/qwen2.5-0.5B.sh"
-
-CKPT_ARGS=(
-   --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/
-   --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/
-)
-
-# num-rollout:100
-ROLLOUT_ARGS=(
-   --prompt-data /root/gsm8k/train.parquet
-   --input-key messages
-   --label-key label
-   --apply-chat-template
-   --rollout-shuffle
-   --rm-type math
-   --num-rollout 500
-   --rollout-batch-size 32
-   --n-samples-per-prompt 8
-   --rollout-max-response-len 1024
-   --rollout-temperature 1
-
-   --global-batch-size 256
-)
-
-EVAL_ARGS=(
-   --eval-interval 20
-   --eval-prompt-data gsm8k /root/gsm8k/test.parquet
-   --n-samples-per-eval-prompt 1
-   --eval-max-response-len 1024
-   --eval-top-k 1
-)
-
-PERF_ARGS=(
-   --tensor-model-parallel-size 1
-   --sequence-parallel
-   --pipeline-model-parallel-size 1
-   --context-parallel-size 1
-   --expert-model-parallel-size 1
-   --expert-tensor-parallel-size 1
-
-   --use-dynamic-batch-size
-   --max-tokens-per-gpu 9216
-)
-
-GRPO_ARGS=(
-   --advantage-estimator grpo
-   --use-kl-loss
-   --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
-)
-
-OPTIMIZER_ARGS=(
-   --optimizer adam
-   --lr 1e-6
-   --lr-decay-style constant
-   --weight-decay 0.1
-   --adam-beta1 0.9
-   --adam-beta2 0.98
-)
-
-WANDB_ARGS=(
-   # --use-wandb
-   # --wandb-host https://wandb.ai/
-   # --wandb-entity samithuang
-   # --wandb-project slime-rl
-   # --wandb-group qwen2.5-0.5B-gsm8k-vllm
-)
-
-VLLM_ARGS=(
-   --rollout-backend vllm
-   --rollout-num-gpus-per-engine 1
-   --rollout-server-concurrency 512
-   --use-slime-router
-   --vllm-gpu-memory-utilization 0.2
-   --slime-router-middleware-paths slime.router.middleware_hub.radix_tree_middleware.RadixTreeMiddleware
-   # Uncomment below to enable colocate mode (IPC weight transfer, same GPU)
-   --colocate
-)
-
-MISC_ARGS=(
-   --attention-dropout 0.0
-   --hidden-dropout 0.0
-   --accumulate-allreduce-grads-in-fp32
-   --attention-softmax-in-fp32
-   --attention-backend flash
-   --deterministic-mode
-)
-
-ray start --head --node-ip-address 127.0.0.1 --num-gpus 2 --disable-usage-stats
-
-ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json='{
-     "env_vars": {
-        "PYTHONPATH": "/root/Megatron-LM",
-        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-        "NCCL_ALGO": "Ring",
-        "NCCL_IB_DISABLE": "1",
-        "NCCL_P2P_DISABLE": "1",
-        "NCCL_SHM_DISABLE": "1",
-        "NCCL_NET_GDR_LEVEL": "0",
-        "NCCL_DEBUG": "INFO",
-        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
-        "CUBLAS_WORKSPACE_CONFIG": ":4096:8"
-     }
-   }' \
-   -- python3 train.py \
-   --actor-num-nodes 1 \
-   --actor-num-gpus-per-node 1 \
-   --num-gpus-per-node 2 \
-   --rollout-num-gpus 1 \
-   --calculate-per-token-loss \
-   ${MODEL_ARGS[@]} \
-   ${CKPT_ARGS[@]} \
-   ${ROLLOUT_ARGS[@]} \
-   ${OPTIMIZER_ARGS[@]} \
-   ${GRPO_ARGS[@]} \
-   ${WANDB_ARGS[@]} \
-   ${PERF_ARGS[@]} \
-   ${EVAL_ARGS[@]} \
-   ${VLLM_ARGS[@]} \
-   ${MISC_ARGS[@]}
diff --git a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh
index 16bc65eaee..c41ea3df82 100644
--- a/scripts/low_precision/run-kimi-k2-Thinking-int4.sh
+++ b/scripts/low_precision/run-kimi-k2-Thinking-int4.sh
@@ -134,7 +134,7 @@ SGLANG_ARGS=(
    #--sglang-deepep-mode auto
 
    # make every dp rank has 128 concurrency
-   --rollout-server-concurrency 1024
+   --sglang-server-concurrency 1024
    --use-slime-router
 )
 
diff --git a/scripts/low_precision/run-qwen3-235B-A22B-int4.sh b/scripts/low_precision/run-qwen3-235B-A22B-int4.sh
index b5ddc7587c..597838f697 100644
--- a/scripts/low_precision/run-qwen3-235B-A22B-int4.sh
+++ b/scripts/low_precision/run-qwen3-235B-A22B-int4.sh
@@ -119,6 +119,7 @@ SGLANG_ARGS=(
   #  --sglang-dp-size 4
    --sglang-ep-size 8
    --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
+   --use-slime-router
 )
 
 
diff --git a/scripts/low_precision/run-qwen3-30B-A3B-int4.sh b/scripts/low_precision/run-qwen3-30B-A3B-int4.sh
index b591047e10..01b4242866 100644
--- a/scripts/low_precision/run-qwen3-30B-A3B-int4.sh
+++ b/scripts/low_precision/run-qwen3-30B-A3B-int4.sh
@@ -114,6 +114,7 @@ SGLANG_ARGS=(
    --rollout-num-gpus-per-engine 1
    --sglang-mem-fraction-static 0.7
    --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
+   --use-slime-router
 )
 
 MISC_ARGS=(
diff --git a/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh b/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh
index 766e0dcc62..fbda2572d2 100644
--- a/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh
+++ b/scripts/low_precision/run-qwen3-30b-a3b-fp8.sh
@@ -127,6 +127,7 @@ SGLANG_ARGS=(
    --sglang-mem-fraction-static 0.6
    --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
    --sglang-expert-parallel-size 8
+   --use-slime-router
    # --use-rollout-routing-replay
 )
 
diff --git a/scripts/models/deepseek-v3.sh b/scripts/models/deepseek-v3.sh
index 8c50d2c940..7efe0e99df 100644
--- a/scripts/models/deepseek-v3.sh
+++ b/scripts/models/deepseek-v3.sh
@@ -43,7 +43,7 @@ MODEL_ARGS=(
 
     # moe
     --num-experts 256
-    --moe-layer-freq "$MOE_LAYER_FREQ"
+    --moe-layer-freq $MOE_LAYER_FREQ
     --moe-ffn-hidden-size 2048
     --moe-router-topk 8
     --moe-shared-expert-intermediate-size 2048
diff --git a/scripts/models/kimi-k2-thinking.sh b/scripts/models/kimi-k2-thinking.sh
index b7fceda591..51fe4fe37a 100644
--- a/scripts/models/kimi-k2-thinking.sh
+++ b/scripts/models/kimi-k2-thinking.sh
@@ -43,7 +43,7 @@ MODEL_ARGS=(
 
     # moe
     --num-experts 384
-    --moe-layer-freq "$MOE_LAYER_FREQ"
+    --moe-layer-freq $MOE_LAYER_FREQ
     --moe-ffn-hidden-size 2048
     --moe-router-topk 8
     --moe-shared-expert-intermediate-size 2048
diff --git a/scripts/models/kimi-k2.sh b/scripts/models/kimi-k2.sh
index eafb7acadd..d9fb1564e3 100644
--- a/scripts/models/kimi-k2.sh
+++ b/scripts/models/kimi-k2.sh
@@ -43,7 +43,7 @@ MODEL_ARGS=(
 
     # moe
     --num-experts 384
-    --moe-layer-freq "$MOE_LAYER_FREQ"
+    --moe-layer-freq $MOE_LAYER_FREQ
     --moe-ffn-hidden-size 2048
     --moe-router-topk 8
     --moe-shared-expert-intermediate-size 2048
diff --git a/scripts/models/moonlight.sh b/scripts/models/moonlight.sh
index bcce99892a..e1e4596f64 100644
--- a/scripts/models/moonlight.sh
+++ b/scripts/models/moonlight.sh
@@ -48,7 +48,7 @@ MODEL_ARGS=(
 
     # moe
     --num-experts 64
-    --moe-layer-freq "$MOE_LAYER_FREQ"
+    --moe-layer-freq $MOE_LAYER_FREQ
     --moe-ffn-hidden-size $MOE_FFN_HIDDEN
     --moe-router-topk 6
     --moe-shared-expert-intermediate-size $MOE_SHARED_EXPERT_INTERMEDIATE_SIZE
diff --git a/scripts/models/qwen3-235B-A22B.sh b/scripts/models/qwen3-235B-A22B.sh
index 1f66355265..32a98679f7 100644
--- a/scripts/models/qwen3-235B-A22B.sh
+++ b/scripts/models/qwen3-235B-A22B.sh
@@ -39,7 +39,7 @@ MODEL_ARGS=(
    --moe-router-score-function softmax
    --moe-token-dispatcher-type alltoall
    --moe-router-topk 8
-   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --moe-layer-freq $MOE_LAYER_FREQ
    --num-experts 128
    --moe-grouped-gemm
    --moe-token-drop-policy probs
diff --git a/scripts/models/qwen3-30B-A3B.sh b/scripts/models/qwen3-30B-A3B.sh
index dfab8682ab..624623d8fe 100644
--- a/scripts/models/qwen3-30B-A3B.sh
+++ b/scripts/models/qwen3-30B-A3B.sh
@@ -39,7 +39,7 @@ MODEL_ARGS=(
    --moe-router-score-function softmax
    --moe-token-dispatcher-type alltoall
    --moe-router-topk 8
-   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --moe-layer-freq $MOE_LAYER_FREQ
    --num-experts 128
    --moe-grouped-gemm
    --moe-token-drop-policy probs
diff --git a/scripts/models/qwen3-next-80B-A3B.sh b/scripts/models/qwen3-next-80B-A3B.sh
index 74e4e02ed8..52650f2cf8 100644
--- a/scripts/models/qwen3-next-80B-A3B.sh
+++ b/scripts/models/qwen3-next-80B-A3B.sh
@@ -44,7 +44,7 @@ MODEL_ARGS=(
    --moe-router-score-function softmax
    --moe-token-dispatcher-type alltoall
    --moe-router-topk 10
-   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --moe-layer-freq $MOE_LAYER_FREQ
    --num-experts 512
    --moe-grouped-gemm
    --moe-token-drop-policy probs
diff --git a/scripts/models/qwen3.5-0.8B.sh b/scripts/models/qwen3.5-0.8B.sh
deleted file mode 100644
index c235c4c1f6..0000000000
--- a/scripts/models/qwen3.5-0.8B.sh
+++ /dev/null
@@ -1,28 +0,0 @@
-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/scripts/models/qwen3.5-35B-A3B.sh b/scripts/models/qwen3.5-35B-A3B.sh
index c3cb7219bb..e96fab0b11 100644
--- a/scripts/models/qwen3.5-35B-A3B.sh
+++ b/scripts/models/qwen3.5-35B-A3B.sh
@@ -44,7 +44,7 @@ MODEL_ARGS=(
    --moe-router-score-function softmax
    --moe-token-dispatcher-type alltoall
    --moe-router-topk 8
-   --moe-layer-freq "$MOE_LAYER_FREQ"
+   --moe-layer-freq $MOE_LAYER_FREQ
    --num-experts 256
    --moe-grouped-gemm
    --moe-token-drop-policy probs
diff --git a/scripts/models/qwen3.5-4B.sh b/scripts/models/qwen3.5-4B.sh
deleted file mode 100644
index 18bea47c9c..0000000000
--- a/scripts/models/qwen3.5-4B.sh
+++ /dev/null
@@ -1,27 +0,0 @@
-MODEL_ARGS=(
-   --spec "slime_plugins.models.qwen3_5" "get_qwen3_5_spec"
-
-   --disable-bias-linear
-   --qk-layernorm
-   --group-query-attention
-   --num-attention-heads 16
-   --num-query-groups 4
-   --kv-channels 256
-   --num-layers 32
-   --hidden-size 2560
-   --ffn-hidden-size 9216
-   --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/scripts/run-deepseek-r1.sh b/scripts/run-deepseek-r1.sh
index 8eb6775412..c307d110ec 100644
--- a/scripts/run-deepseek-r1.sh
+++ b/scripts/run-deepseek-r1.sh
@@ -121,18 +121,12 @@ SGLANG_ARGS=(
    --sglang-moe-dense-tp-size 1
    --sglang-enable-dp-lm-head
 
-    # enable deepep for sglang
-    --sglang-moe-a2a-backend deepep
-    --sglang-deepep-mode auto
+   # enable deepep for sglang
+   --sglang-moe-a2a-backend deepep
+   --sglang-deepep-mode auto
 
    # make every dp rank has 128 concurrency
-   --rollout-server-concurrency 1024
-    # mtp
-    --sglang-speculative-algorithm EAGLE
-    --sglang-speculative-num-steps 3
-    --sglang-speculative-eagle-topk 1
-    --sglang-speculative-num-draft-tokens 4
-
+   --sglang-server-concurrency 1024
 )
 
 MISC_ARGS=(
diff --git a/scripts/run-glm4.7-355B-A32B.sh b/scripts/run-glm4.5-355B-A32B.sh
similarity index 60%
rename from scripts/run-glm4.7-355B-A32B.sh
rename to scripts/run-glm4.5-355B-A32B.sh
index aea94f4b6d..04b603fc24 100644
--- a/scripts/run-glm4.7-355B-A32B.sh
+++ b/scripts/run-glm4.5-355B-A32B.sh
@@ -15,8 +15,6 @@ set -ex
 # will prevent ray from buffering stdout/stderr
 export PYTHONBUFFERED=16
 
-unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
-
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
     HAS_NVLINK=1
@@ -29,8 +27,8 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
 source "${SCRIPT_DIR}/models/glm4.5-355B-A32B.sh"
 
 CKPT_ARGS=(
-   --hf-checkpoint $BASE_DIR/GLM-4.7-355B-A32B
-   --ref-load $BASE_DIR/GLM-4.7-355B-A32B_torch_dist/
+   --hf-checkpoint $BASE_DIR/GLM-4.5-355B-A32B
+   --ref-load $BASE_DIR/GLM-4.5-355B-A32B_torch_dist/
 )
 
 ROLLOUT_ARGS=(
@@ -107,7 +105,7 @@ OPTIMIZER_ARGS=(
 WANDB_ARGS=(
    # --use-wandb
    # --wandb-project slime-dev
-   # --wandb-group glm4.7-355B
+   # --wandb-group qwen3-235B-sft
 )
 
 SGLANG_ARGS=(
@@ -121,9 +119,9 @@ SGLANG_ARGS=(
 
    # mtp
    --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
+   --sglang-speculative-num-steps 2
    --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --sglang-speculative-num-draft-tokens 3
 
 )
 
@@ -141,54 +139,56 @@ MISC_ARGS=(
    --moe-enable-deepep
 )
 
-ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-8}
-ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8}
-SOCKET_IFNAME=${SOCKET_IFNAME:-eth0}
-
-MASTER_ADDR=${MASTER_ADDR:-}
-if [ -z "${MASTER_ADDR}" ]; then
-  echo "MASTER_ADDR is not set. Please set it to the master node address."
-  exit 1
-fi
-
-# launch the master node of ray
+# launch the master node of ray in container
+export MASTER_ADDR=${MLP_WORKER_0_HOST}
 export no_proxy="127.0.0.1,${MASTER_ADDR}"
-ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-HOSTFILE=${HOSTFILE:-}
-if [ -n "${HOSTFILE}" ]; then
-  for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do
-    if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then
-      continue
-    fi
-    echo "Starting Ray worker on ${WORKER_IP}"
-    ssh root@"${WORKER_IP}" \
-      "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" &
-  done
-  wait
-fi
-
-RUNTIME_ENV_JSON=$(cat <<EOF_JSON
-{
-  "env_vars": {
-    "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}",
-    "GLOO_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "TP_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "MASTER_ADDR": "${MASTER_ADDR}",
-    "PYTHONPATH": "/root/Megatron-LM/",
-    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-    "NCCL_NVLS_ENABLE": "${HAS_NVLINK}"
-  }
-}
-EOF_JSON
-)
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do
+  if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" &
+done
+wait
 
 ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --runtime-env-json='{
+     "env_vars": {
+        "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}",
+        "GLOO_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}",
+        "TP_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}",
+        "MASTER_ADDR": "${MLP_WORKER_0_HOST}",
+        "PYTHONPATH": "/root/Megatron-LM/",
+        "NCCL_CUMEM_ENABLE": "0",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": "20",
+        "NCCL_IB_TC": "160",
+        "NCCL_PXN_DISABLE": "0",
+        "NCCL_IB_GID_INDEX": "3",
+        "NCCL_NET_GDR_LEVEL": "4",
+        "NCCL_IB_RETRY_CNT": "7",
+        "NCCL_IB_TIMEOUT": "32",
+        "NCCL_IB_QPS_PER_CONNECTION": "8",
+        "NCCL_P2P_LEVEL": "NVL",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
+        "NCCL_NVLS_ENABLE": "0",
+        "NCCL_MIN_CTAS": "4",
+        "OMPI_MCA_pml": "ob1",
+        "OMPI_MCA_btl": "^openib",
+        "OMPI_MCA_routed": "direct",
+        "OMPI_MCA_routed_radix": "1024",
+        "OMPI_MCA_plm_rsh_no_tree_spawn": "1",
+        "OMPI_MCA_oob_tcp_if_include": "${MLP_SOCKET_IFNAME}",
+        "OMPI_MCA_btl_tcp_if_include": "${MLP_SOCKET_IFNAME}"
+     }
+   }' \
    -- python3 train.py \
-   --actor-num-nodes "${ACTOR_NUM_NODES}" \
-   --actor-num-gpus-per-node "${ACTOR_NUM_GPUS_PER_NODE}" \
+   --actor-num-nodes 8 \
+   --actor-num-gpus-per-node 8 \
    --colocate \
+   --save-debug-rollout-data data.pt \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
diff --git a/scripts/run-glm4.7-30B-A3B-8gpus.sh b/scripts/run-glm4.7-30B-A3B-8gpus.sh
new file mode 100755
index 0000000000..62c56529be
--- /dev/null
+++ b/scripts/run-glm4.7-30B-A3B-8gpus.sh
@@ -0,0 +1,175 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 sglang
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+export PYTHONBUFFERED=16
+
+NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/models/glm4.7-30B-A3B.sh"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/GLM-4.7-Flash
+   --ref-load /root/GLM-4.7-Flash_torch_dist
+   --load /root/GLM-4.7-Flash_slime/
+   --save /root/GLM-4.7-Flash_slime/
+   --save-interval 20
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
+)
+
+PERF_ARGS=(
+   --tensor-model-parallel-size 1
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
+   --expert-model-parallel-size 8
+   --expert-tensor-parallel-size 1
+
+   --recompute-granularity full
+   --recompute-method uniform
+   --recompute-num-layers 1
+
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 8192
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
+)
+
+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
+)
+
+WANDB_ARGS=(
+   #--use-wandb
+   # --wandb-project slime-dev
+   # --wandb-group glm4.7-flash
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.7
+   --sglang-enable-dp-attention
+   --sglang-dp-size 8
+   --sglang-enable-dp-lm-head
+   --sglang-moe-dense-tp-size 1
+
+   # MTP speculative decoding (EAGLE) — speeds up inference
+   # Uncomment the following lines to enable:
+   # --sglang-speculative-algorithm EAGLE
+   # --sglang-speculative-num-steps 2
+   # --sglang-speculative-eagle-topk 1
+   # --sglang-speculative-num-draft-tokens 3
+
+   --sglang-cuda-graph-max-bs 16
+   --sglang-max-running-requests 64
+)
+
+MISC_ARGS=(
+   # default dropout in megatron is 0.1
+   --attention-dropout 0.0
+   --hidden-dropout 0.0
+   # should be good for model performance
+   --accumulate-allreduce-grads-in-fp32
+   --attention-softmax-in-fp32
+   # need to comment this when using model with MLA
+   --attention-backend flash
+
+   --moe-token-dispatcher-type alltoall
+)
+
+# MTP training — uncomment to enable MTP layer training
+# Requires --mtp-num-layers in MODEL_ARGS (add to glm4.7-30B-A3B.sh or override below)
+# MODEL_ARGS+=(--mtp-num-layers 1)
+SPEC_ARGS=(
+   # --enable-mtp-training
+   # --mtp-loss-scaling-factor 0.2
+)
+
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+  }
+}"
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train.py \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
+   --colocate \
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${SPEC_ARGS[@]}
diff --git a/scripts/run-glm4.7-30B-A3B.sh b/scripts/run-glm4.7-30B-A3B.sh
index 1b52b25460..11f4f20648 100644
--- a/scripts/run-glm4.7-30B-A3B.sh
+++ b/scripts/run-glm4.7-30B-A3B.sh
@@ -15,8 +15,6 @@ set -ex
 # will prevent ray from buffering stdout/stderr
 export PYTHONBUFFERED=16
 
-unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
-
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
     HAS_NVLINK=1
@@ -44,6 +42,7 @@ ROLLOUT_ARGS=(
 
    --num-rollout 3000
    --rollout-batch-size 128
+   #--over-sampling-batch-size 256
    --n-samples-per-prompt 8
    --rollout-max-response-len 32768
    --rollout-temperature 1.0
@@ -56,13 +55,13 @@ EVAL_ARGS=(
    --eval-interval 20
    --eval-prompt-data aime24 $BASE_DIR/rl_data/aime-2024.jsonl
    --n-samples-per-eval-prompt 2
-   --eval-max-response-len 32768
-   --eval-temperature 1.0
+   --eval-max-response-len 16384
+   --eval-temperature 0.6
    --eval-top-p 0.95
 )
 
 PERF_ARGS=(
-   --tensor-model-parallel-size 2
+   --tensor-model-parallel-size 4
    --sequence-parallel
    --pipeline-model-parallel-size 2
    --context-parallel-size 2
@@ -78,12 +77,6 @@ PERF_ARGS=(
    --max-tokens-per-gpu 32768
 )
 
-MTP_ARGS=(
-   --mtp-num-layers 1
-   --enable-mtp-training
-   --mtp-loss-scaling-factor 0.2
-)
-
 GRPO_ARGS=(
    --advantage-estimator grpo
    --use-kl-loss
@@ -122,9 +115,9 @@ SGLANG_ARGS=(
 
    # mtp
    --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
+   --sglang-speculative-num-steps 2
    --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --sglang-speculative-num-draft-tokens 3
 
    --sglang-cuda-graph-max-bs 64
 
@@ -146,24 +139,56 @@ MISC_ARGS=(
 )
 
 # launch the master node of ray in container
-export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
-ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-# Build the runtime environment JSON with proper variable substitution
-RUNTIME_ENV_JSON="{
-  \"env_vars\": {
-    \"PYTHONPATH\": \"/root/Megatron-LM/\",
-    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
-  }
-}"
+export MASTER_ADDR=${MLP_WORKER_0_HOST}
+export no_proxy="127.0.0.1,${MASTER_ADDR}"
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
+
+for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do
+  if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats" &
+done
+wait
 
 ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   --runtime-env-json='{
+     "env_vars": {
+        "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}",
+        "GLOO_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}",
+        "TP_SOCKET_IFNAME": "${MLP_SOCKET_IFNAME}",
+        "MASTER_ADDR": "${MLP_WORKER_0_HOST}",
+        "PYTHONPATH": "/root/Megatron-LM/",
+        "NCCL_CUMEM_ENABLE": "0",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+        "NVTE_BWD_LAYERNORM_SM_MARGIN": "20",
+        "NCCL_IB_TC": "160",
+        "NCCL_PXN_DISABLE": "0",
+        "NCCL_IB_GID_INDEX": "3",
+        "NCCL_NET_GDR_LEVEL": "4",
+        "NCCL_IB_RETRY_CNT": "7",
+        "NCCL_IB_TIMEOUT": "32",
+        "NCCL_IB_QPS_PER_CONNECTION": "8",
+        "NCCL_P2P_LEVEL": "NVL",
+        "TORCH_NCCL_AVOID_RECORD_STREAMS": "1",
+        "NCCL_NVLS_ENABLE": "0",
+        "NCCL_MIN_CTAS": "4",
+        "OMPI_MCA_pml": "ob1",
+        "OMPI_MCA_btl": "^openib",
+        "OMPI_MCA_routed": "direct",
+        "OMPI_MCA_routed_radix": "1024",
+        "OMPI_MCA_plm_rsh_no_tree_spawn": "1",
+        "OMPI_MCA_oob_tcp_if_include": "${MLP_SOCKET_IFNAME}",
+        "OMPI_MCA_btl_tcp_if_include": "${MLP_SOCKET_IFNAME}"
+     }
+   }' \
    -- python3 train.py \
-   --actor-num-nodes 1 \
+   --actor-num-nodes 2 \
    --actor-num-gpus-per-node 8 \
    --colocate \
+   --save-debug-rollout-data "data.pt" \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
@@ -173,5 +198,4 @@ ray job submit --address="http://127.0.0.1:8265" \
    ${PERF_ARGS[@]} \
    ${EVAL_ARGS[@]} \
    ${SGLANG_ARGS[@]} \
-   ${MISC_ARGS[@]} \
-   ${MTP_ARGS[@]}
+   ${MISC_ARGS[@]}
diff --git a/scripts/run-kimi-k2-Instruct.sh b/scripts/run-kimi-k2-Instruct.sh
index 0c9fc2c3d4..3a591b923a 100644
--- a/scripts/run-kimi-k2-Instruct.sh
+++ b/scripts/run-kimi-k2-Instruct.sh
@@ -132,7 +132,7 @@ SGLANG_ARGS=(
 #    --sglang-deepep-mode auto
 
    # make every dp rank has 128 concurrency
-   --rollout-server-concurrency 1024
+   --sglang-server-concurrency 1024
 )
 
 
diff --git a/scripts/run-kimi-k2-Thinking.sh b/scripts/run-kimi-k2-Thinking.sh
index 533019d697..25cc3c475f 100644
--- a/scripts/run-kimi-k2-Thinking.sh
+++ b/scripts/run-kimi-k2-Thinking.sh
@@ -134,7 +134,7 @@ SGLANG_ARGS=(
    # --sglang-deepep-mode auto
 
    # make every dp rank has 128 concurrency
-   --rollout-server-concurrency 1024
+   --sglang-server-concurrency 1024
 )
 
 
diff --git a/scripts/run-qwen2.5-0.5B-gb10-smoke.sh b/scripts/run-qwen2.5-0.5B-gb10-smoke.sh
deleted file mode 100755
index 294c6836fc..0000000000
--- a/scripts/run-qwen2.5-0.5B-gb10-smoke.sh
+++ /dev/null
@@ -1,113 +0,0 @@
-#!/bin/bash
-# Minimal GRPO smoke test for slime on NVIDIA DGX Spark (GB10, single GPU).
-# Goal: exercise the full rollout → reward → policy-update loop for one tiny
-# step and exit cleanly. Used only to validate the GB10 port; not a training
-# recipe.
-#
-# Prerequisites:
-#   - /root/Qwen2.5-0.5B-Instruct                    (HF checkpoint)
-#   - /root/Qwen2.5-0.5B-Instruct_torch_dist         (from tools/convert_hf_to_torch_dist.py)
-#   - /root/dapo-math-17k/dapo-math-17k.jsonl
-
-set -ex
-
-# clean any leftover ray/sglang
-pkill -9 sglang 2>/dev/null || true
-ray stop --force 2>/dev/null || true
-pkill -9 ray python 2>/dev/null || true
-sleep 2
-
-export PYTHONBUFFERED=1
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
-source "${SCRIPT_DIR}/models/qwen2.5-0.5B.sh"
-
-CKPT_ARGS=(
-   --hf-checkpoint /root/Qwen2.5-0.5B-Instruct/
-   --ref-load /root/Qwen2.5-0.5B-Instruct_torch_dist/
-   --save /tmp/slime_smoke_save/
-   --save-interval 9999
-)
-
-ROLLOUT_ARGS=(
-   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
-   --input-key prompt
-   --label-key label
-   --apply-chat-template
-   --rollout-shuffle
-   --rm-type deepscaler
-
-   --num-rollout 1
-   --rollout-batch-size 2
-   --n-samples-per-prompt 2
-   --num-steps-per-rollout 1
-   --global-batch-size 4
-
-   --rollout-max-response-len 256
-   --rollout-temperature 1
-)
-
-PERF_ARGS=(
-   --tensor-model-parallel-size 1
-   --pipeline-model-parallel-size 1
-   --context-parallel-size 1
-   --expert-model-parallel-size 1
-   --expert-tensor-parallel-size 1
-
-   --use-dynamic-batch-size
-   --max-tokens-per-gpu 2048
-)
-
-GRPO_ARGS=(
-   --advantage-estimator grpo
-   --use-kl-loss
-   --kl-loss-coef 0.00
-   --kl-loss-type low_var_kl
-   --entropy-coef 0.00
-   --eps-clip 0.2
-   --eps-clip-high 0.28
-)
-
-OPTIMIZER_ARGS=(
-   --optimizer adam
-   --lr 1e-6
-   --lr-decay-style constant
-   --weight-decay 0.1
-   --adam-beta1 0.9
-   --adam-beta2 0.98
-)
-
-SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 1
-   --sglang-mem-fraction-static 0.4
-)
-
-MISC_ARGS=(
-   --attention-dropout 0.0
-   --hidden-dropout 0.0
-   --accumulate-allreduce-grads-in-fp32
-   --attention-softmax-in-fp32
-   --attention-backend flash
-)
-
-ray start --head --node-ip-address 127.0.0.1 --num-gpus 1 --disable-usage-stats
-
-ray job submit --address="http://127.0.0.1:8265" \
-   --runtime-env-json='{
-     "env_vars": {
-        "PYTHONPATH": "/root/src/Megatron-LM",
-        "CUDA_DEVICE_MAX_CONNECTIONS": "1"
-     }
-   }' \
-   -- python3 train.py \
-   --actor-num-nodes 1 \
-   --actor-num-gpus-per-node 1 \
-   --colocate \
-   "${MODEL_ARGS[@]}" \
-   "${CKPT_ARGS[@]}" \
-   "${ROLLOUT_ARGS[@]}" \
-   "${OPTIMIZER_ARGS[@]}" \
-   "${GRPO_ARGS[@]}" \
-   "${PERF_ARGS[@]}" \
-   "${SGLANG_ARGS[@]}" \
-   "${MISC_ARGS[@]}"
diff --git a/scripts/run-qwen2.5-0.5B-reproducibility.sh b/scripts/run-qwen2.5-0.5B-reproducibility.sh
index d1d2d3bf1c..695a282d74 100644
--- a/scripts/run-qwen2.5-0.5B-reproducibility.sh
+++ b/scripts/run-qwen2.5-0.5B-reproducibility.sh
@@ -126,6 +126,7 @@ ray job submit --address="http://127.0.0.1:8265" \
    --actor-num-gpus-per-node 8 \
    --colocate \
    --calculate-per-token-loss \
+   --use-slime-router \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
diff --git a/scripts/run-qwen3-235B-A22B.sh b/scripts/run-qwen3-235B-A22B.sh
index 28ef17a221..c9958c629a 100644
--- a/scripts/run-qwen3-235B-A22B.sh
+++ b/scripts/run-qwen3-235B-A22B.sh
@@ -161,7 +161,7 @@ RUNTIME_ENV_JSON="{
   \"env_vars\": {
     \"PYTHONPATH\": \"/root/Megatron-LM/\",
     \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
-    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
     \"no_proxy\": \"${no_proxy}\",
     \"MASTER_ADDR\": \"${MASTER_ADDR}\"
   }
diff --git a/scripts/run-qwen3-4B-fsdp.sh b/scripts/run-qwen3-4B-fsdp.sh
new file mode 100644
index 0000000000..71e48e21a7
--- /dev/null
+++ b/scripts/run-qwen3-4B-fsdp.sh
@@ -0,0 +1,146 @@
+#!/bin/bash
+
+# for rerun the task
+pkill -9 sglang
+sleep 3
+ray stop --force
+pkill -9 ray
+pkill -9 python
+sleep 3
+pkill -9 ray
+pkill -9 python
+
+
+
+
+set -ex
+
+# will prevent ray from buffering stdout/stderr
+export PYTHONBUFFERED=16
+export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7
+NVLINK_COUNT=$(nvidia-smi | grep -o "NVLink" | wc -l)
+if [ "$NVLINK_COUNT" -gt 0 ]; then
+    HAS_NVLINK=1
+else
+    HAS_NVLINK=0
+fi
+echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
+
+
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+
+RUN_ID=${RUN_ID:-"run_$(date +%Y%m%d_%H%M%S)"}
+LOAD_SAVE_PATH="/root/shared_data/${RUN_ID}/checkpoints"
+
+CKPT_ARGS=(
+   --hf-checkpoint /root/Qwen3-4B
+   --load /root/Qwen3-4B
+   --ref-load /root/Qwen3-4B
+)
+
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
+   --rollout-shuffle
+   --balance-data
+   --rm-type deepscaler
+   --num-rollout 100
+   --rollout-batch-size 8
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 4096
+   --rollout-temperature 1
+   --global-batch-size 64
+)
+
+GRPO_ARGS=(
+   --use-kl-loss
+   --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
+)
+
+OPTIMIZER_ARGS=(
+   --optimizer adam
+   --lr 1e-6
+   --lr-decay-style constant
+   --weight-decay 0.1
+   --adam-beta1 0.9
+   --adam-beta2 0.98
+)
+
+if [ -z "${WANDB_API_KEY}" ]; then
+   WANDB_ARGS=()
+else
+   WANDB_ARGS=(
+      --use-wandb
+      --wandb-project slime-dev-mcore-fsdp
+      --wandb-group qwen3-4B-fsdp-1130-ref
+      --wandb-key "${WANDB_API_KEY}"
+   )
+fi
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 1
+   --sglang-mem-fraction-static 0.75
+   --sglang-decode-log-interval 1000
+   --sglang-chunked-prefill-size 4096
+   --sglang-attention-backend fa3
+)
+
+TRAIN_BACKEND_ARGS=(
+   --train-backend fsdp
+   --update-weight-buffer-size 536870912
+   --gradient-checkpointing
+   --attn-implementation flash_attention_3
+   --train-env-vars '{"PYTORCH_CUDA_ALLOC_CONF":"expandable_segments:True"}'
+)
+
+PERF_ARGS=(
+   --use-dynamic-batch-size
+   --max-tokens-per-gpu 9216
+)
+
+MISC_ARGS=(
+   --actor-num-nodes 1
+   --actor-num-gpus-per-node 8
+   --colocate
+   --use-fault-tolerance
+   --dump-details /root/shared_data/qwen3-4B-fsdp-1116-noref/dump_details
+   # --fsdp-cpu-offload
+)
+
+# launch the master node of ray in container - 8 GPUs for training
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats
+
+
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/:${SCRIPT_DIR}\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\"
+  }
+}"
+
+
+ray job submit --address="http://127.0.0.1:8265" \
+   --runtime-env-json="${RUNTIME_ENV_JSON}" \
+   -- python3 train.py \
+   "${CKPT_ARGS[@]}" \
+   "${ROLLOUT_ARGS[@]}" \
+   "${OPTIMIZER_ARGS[@]}" \
+   "${GRPO_ARGS[@]}" \
+   "${WANDB_ARGS[@]}" \
+   "${SGLANG_ARGS[@]}" \
+   "${TRAIN_BACKEND_ARGS[@]}" \
+   "${PERF_ARGS[@]}" \
+   "${MISC_ARGS[@]}"
+
+
+
diff --git a/scripts/run-qwen3-4B.sh b/scripts/run-qwen3-4B.sh
index 17023da848..ebfe216033 100644
--- a/scripts/run-qwen3-4B.sh
+++ b/scripts/run-qwen3-4B.sh
@@ -23,17 +23,6 @@ else
 fi
 echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)"
 
-if command -v nvidia-smi >/dev/null 2>&1; then
-    DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l | tr -d ' ')
-else
-    DETECTED_GPUS=0
-fi
-NUM_GPUS=${NUM_GPUS:-${DETECTED_GPUS}}
-if [ -z "$NUM_GPUS" ] || [ "$NUM_GPUS" -le 0 ]; then
-    NUM_GPUS=8
-fi
-echo "NUM_GPUS: $NUM_GPUS"
-
 SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
 source "${SCRIPT_DIR}/models/qwen3-4B.sh"
 
@@ -132,7 +121,7 @@ MISC_ARGS=(
 
 # launch the master node of ray in container
 export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
-ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus ${NUM_GPUS} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
 
 # Build the runtime environment JSON with proper variable substitution
 RUNTIME_ENV_JSON="{
@@ -147,7 +136,7 @@ ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
    -- python3 train.py \
    --actor-num-nodes 1 \
-   --actor-num-gpus-per-node ${NUM_GPUS} \
+   --actor-num-gpus-per-node 8 \
    --colocate \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
@@ -158,4 +147,4 @@ ray job submit --address="http://127.0.0.1:8265" \
    ${PERF_ARGS[@]} \
    ${EVAL_ARGS[@]} \
    ${SGLANG_ARGS[@]} \
-   ${MISC_ARGS[@]}
+   ${MISC_ARGS[@]}
\ No newline at end of file
diff --git a/scripts/run-qwen3-next-80B-A3B.sh b/scripts/run-qwen3-next-80B-A3B.sh
index 1b445c534d..8f09158203 100644
--- a/scripts/run-qwen3-next-80B-A3B.sh
+++ b/scripts/run-qwen3-next-80B-A3B.sh
@@ -12,12 +12,12 @@ pkill -9 python
 
 set -ex
 
-if [ -z "${BASE_FOLDER:-}" ]; then
+# if base folder not set raise error
+if [ -z "${BASE_FOLDER}" ]; then
   echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints."
   exit 1
 fi
 
-MASTER_ADDR=${MASTER_ADDR:-}
 if [ -z "${MASTER_ADDR}" ]; then
   echo "MASTER_ADDR is not set. Please set it to the master node address."
   exit 1
@@ -26,14 +26,6 @@ fi
 # will prevent ray from buffering stdout/stderr
 export PYTHONBUFFERED=16
 
-# unset proxy to avoid distributed startup issues
-unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
-
-ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-4}
-ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8}
-CP_SIZE=${CP_SIZE:-4}
-SOCKET_IFNAME=${SOCKET_IFNAME:-eth0}
-
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
     HAS_NVLINK=1
@@ -46,34 +38,35 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
 source "${SCRIPT_DIR}/models/qwen3-next-80B-A3B.sh"
 
 CKPT_ARGS=(
-   --hf-checkpoint "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking"
-   --ref-load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_torch_dist"
-   --load "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/"
-   --save "${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/"
+   --hf-checkpoint ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking
+   --ref-load ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_torch_dist
+   --load ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/
+   --save ${BASE_FOLDER}/Qwen3-Next-80B-A3B-Thinking_slime/
    --save-interval 20
 )
 
 ROLLOUT_ARGS=(
-   --prompt-data "${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl"
+   --prompt-data ${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl
    --input-key prompt
    --label-key label
    --apply-chat-template
    --rollout-shuffle
    --rm-type deepscaler
    --num-rollout 3000
-   --rollout-batch-size 8
+   --rollout-batch-size 32
    --n-samples-per-prompt 8
-   --rollout-max-response-len 32768
+   --rollout-max-response-len 8192
    --rollout-temperature 1
-   --global-batch-size 64
+
+   --global-batch-size 256
    --balance-data
 )
 
 EVAL_ARGS=(
    --eval-interval 20
-   --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl"
-   --n-samples-per-eval-prompt 8
-   --eval-max-response-len 32768
+   --eval-prompt-data aime ${BASE_FOLDER}/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
    --eval-top-p 1
 )
 
@@ -81,8 +74,7 @@ PERF_ARGS=(
    --tensor-model-parallel-size 2
    --sequence-parallel
    --pipeline-model-parallel-size 4
-   --decoder-last-pipeline-num-layers 9
-   --context-parallel-size "${CP_SIZE}"
+   --context-parallel-size 2
    --expert-model-parallel-size 8
    --expert-tensor-parallel-size 1
 
@@ -90,12 +82,14 @@ PERF_ARGS=(
    --recompute-method uniform
    --recompute-num-layers 1
 
+   # --micro-batch-size 1
    --use-dynamic-batch-size
    --max-tokens-per-gpu 8192
 )
 
 GRPO_ARGS=(
    --advantage-estimator gspo
+   #--use-kl-loss
    --kl-loss-coef 0.00
    --kl-loss-type low_var_kl
    --kl-coef 0.00
@@ -117,9 +111,9 @@ OPTIMIZER_ARGS=(
 )
 
 WANDB_ARGS=(
-   # --use-wandb
+   #--use-wandb
    # --wandb-project slime-dev
-   # --wandb-group qwen3-next-80B-A3B-32k
+   # --wandb-group qwen3-next-80B-A3B-test
    # --wandb-key ${WANDB_KEY}
 )
 
@@ -127,71 +121,76 @@ SGLANG_ARGS=(
    --rollout-num-gpus-per-engine 8
    --sglang-mem-fraction-static 0.8
    --sglang-ep-size 8
+   
    --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 128)
 
    # mtp
    --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
+   --sglang-speculative-num-steps 2
    --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --sglang-speculative-num-draft-tokens 3
 
-   --sglang-max-running-requests 256
+   --sglang-max-running-requests 512
 )
 
 MISC_ARGS=(
+   # default dropout in megatron is 0.1
    --attention-dropout 0.0
    --hidden-dropout 0.0
+   # should be good for model performance
    --accumulate-allreduce-grads-in-fp32
    --attention-softmax-in-fp32
+   # need to comment this when using model with MLA
    --attention-backend flash
+
    --moe-token-dispatcher-type flex
    --moe-enable-deepep
 )
 
-export no_proxy="127.0.0.1,${MASTER_ADDR}"
-ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-HOSTFILE=${HOSTFILE:-}
-if [ -n "${HOSTFILE}" ]; then
-  for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do
-    if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then
-      continue
-    fi
-    echo "Starting Ray worker on ${WORKER_IP}"
-    ssh root@"${WORKER_IP}" \
-      "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" &
-  done
-  wait
-fi
+SPEC_ARGS={
+   # --mtp-num-layers 1
+   # --enable-mtp-training
+   # --mtp-loss-scaling-factor 0.2
+}
 
-RUNTIME_ENV_JSON=$(cat <<EOF_JSON
-{
-  "env_vars": {
-    "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}",
-    "GLOO_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "TP_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "MASTER_ADDR": "${MASTER_ADDR}",
-    "PYTHONPATH": "/root/Megatron-LM/",
-    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-    "NCCL_NVLS_ENABLE": "${HAS_NVLINK}"
+# launch the master node of ray in container
+export no_proxy="127.0.0.1,${MASTER_ADDR}"
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do
+  if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then
+    continue
+  fi
+  echo "Starting Ray worker on ${WORKER_IP}"
+  ssh root@"${WORKER_IP}" \
+    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" &
+done
+wait
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\"
+    \"no_proxy\": \"${no_proxy}\",
+    \"MASTER_ADDR\": \"${MASTER_ADDR}\"
   }
-}
-EOF_JSON
-)
+}"
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
    -- python3 train.py \
-   --actor-num-nodes "${ACTOR_NUM_NODES}" \
-   --actor-num-gpus-per-node "${ACTOR_NUM_GPUS_PER_NODE}" \
+   --actor-num-nodes 4 \
+   --actor-num-gpus-per-node 8 \
    --colocate \
-   "${MODEL_ARGS[@]}" \
-   "${CKPT_ARGS[@]}" \
-   "${ROLLOUT_ARGS[@]}" \
-   "${OPTIMIZER_ARGS[@]}" \
-   "${GRPO_ARGS[@]}" \
-   "${WANDB_ARGS[@]}" \
-   "${PERF_ARGS[@]}" \
-   "${EVAL_ARGS[@]}" \
-   "${SGLANG_ARGS[@]}" \
-   "${MISC_ARGS[@]}"
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]} \
+   ${SPEC_ARGS[@]}
diff --git a/scripts/run-qwen3.5-27B.sh b/scripts/run-qwen3.5-27B.sh
index 6c31215e8a..f0abd623ae 100755
--- a/scripts/run-qwen3.5-27B.sh
+++ b/scripts/run-qwen3.5-27B.sh
@@ -15,25 +15,9 @@ set -ex
 # will prevent ray from buffering stdout/stderr
 export PYTHONBUFFERED=16
 
-# unset proxy to avoid distributed startup issues
+# unset proxy to avoid issues
 unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
 
-if [ -z "${BASE_FOLDER:-}" ]; then
-  echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints."
-  exit 1
-fi
-
-MASTER_ADDR=${MASTER_ADDR:-}
-if [ -z "${MASTER_ADDR}" ]; then
-  echo "MASTER_ADDR is not set. Please set it to the master node address."
-  exit 1
-fi
-
-ACTOR_NUM_NODES=${ACTOR_NUM_NODES:-4}
-ACTOR_NUM_GPUS_PER_NODE=${ACTOR_NUM_GPUS_PER_NODE:-8}
-CP_SIZE=${CP_SIZE:-4}
-SOCKET_IFNAME=${SOCKET_IFNAME:-eth0}
-
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
     HAS_NVLINK=1
@@ -46,43 +30,43 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
 source "${SCRIPT_DIR}/models/qwen3.5-27B.sh"
 
 CKPT_ARGS=(
-   --hf-checkpoint "${BASE_FOLDER}/Qwen3.5-27B"
-   --ref-load "${BASE_FOLDER}/Qwen3.5-27B_torch_dist/"
-   --load "${BASE_FOLDER}/Qwen3.5-27B_slime/"
-   --save "${BASE_FOLDER}/Qwen3.5-27B_slime/"
+   --hf-checkpoint /root/Qwen3.5-27B
+   --ref-load /root/Qwen3.5-27B_torch_dist/
+   --load /root/Qwen3.5-27B_slime
+   --save /root/Qwen3.5-27B_slime
    --save-interval 20
 )
 
 ROLLOUT_ARGS=(
-   --prompt-data "${BASE_FOLDER}/dapo-math-17k/dapo-math-17k.jsonl"
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
    --input-key prompt
    --label-key label
    --apply-chat-template
    --rollout-shuffle
    --rm-type deepscaler
    --num-rollout 3000
-   --rollout-batch-size 8
+   --rollout-batch-size 32
    --n-samples-per-prompt 8
-   --rollout-max-response-len 32768
-   --rollout-temperature 1.0
-   --global-batch-size 64
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 256
    --balance-data
 )
 
 EVAL_ARGS=(
    --eval-interval 20
-   --eval-prompt-data aime "${BASE_FOLDER}/aime-2024/aime-2024.jsonl"
-   --n-samples-per-eval-prompt 8
-   --eval-max-response-len 32768
+   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
    --eval-top-p 1
 )
 
 PERF_ARGS=(
    --tensor-model-parallel-size 4
    --sequence-parallel
-   --pipeline-model-parallel-size 2
-   --decoder-last-pipeline-num-layers 30
-   --context-parallel-size "${CP_SIZE}"
+   --pipeline-model-parallel-size 1
+   --context-parallel-size 1
    --expert-model-parallel-size 1
    --expert-tensor-parallel-size 1
 
@@ -90,18 +74,19 @@ PERF_ARGS=(
    --recompute-method uniform
    --recompute-num-layers 1
 
+   # --micro-batch-size 1
    --use-dynamic-batch-size
-   --calculate-per-token-loss
-   --max-tokens-per-gpu 8192
+   --max-tokens-per-gpu 20480
 )
 
 GRPO_ARGS=(
    --advantage-estimator grpo
+   --use-kl-loss
    --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
 )
 
 OPTIMIZER_ARGS=(
@@ -111,79 +96,64 @@ OPTIMIZER_ARGS=(
    --weight-decay 0.1
    --adam-beta1 0.9
    --adam-beta2 0.98
+
    --optimizer-cpu-offload
    --overlap-cpu-optimizer-d2h-h2d
    --use-precision-aware-optimizer
 )
 
 WANDB_ARGS=(
-   # --use-wandb
+   #--use-wandb
    # --wandb-project slime-dev
-   # --wandb-group qwen3.5-27B-32k
+   # --wandb-group qwen3.5-27B-test
    # --wandb-key ${WANDB_KEY}
 )
 
 SGLANG_ARGS=(
-   --rollout-num-gpus-per-engine 2
-   --sglang-mem-fraction-static 0.75
-   --sglang-speculative-algorithm EAGLE
-   --sglang-speculative-num-steps 3
-   --sglang-speculative-eagle-topk 1
-   --sglang-speculative-num-draft-tokens 4
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.7
+   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
 )
 
 MISC_ARGS=(
+   # default dropout in megatron is 0.1
    --attention-dropout 0.0
    --hidden-dropout 0.0
+   # should be good for model performance
    --accumulate-allreduce-grads-in-fp32
    --attention-softmax-in-fp32
+   # need to comment this when using model with MLA
    --attention-backend flash
 )
 
+# launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
 export no_proxy="127.0.0.1,${MASTER_ADDR}"
-ray start --head --node-ip-address "${MASTER_ADDR}" --num-gpus "${ACTOR_NUM_GPUS_PER_NODE}" --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-
-HOSTFILE=${HOSTFILE:-}
-if [ -n "${HOSTFILE}" ]; then
-  for WORKER_IP in $(awk '{print $1}' "${HOSTFILE}"); do
-    if [[ "${WORKER_IP}" == "${MASTER_ADDR}" ]]; then
-      continue
-    fi
-    echo "Starting Ray worker on ${WORKER_IP}"
-    ssh root@"${WORKER_IP}" \
-      "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus ${ACTOR_NUM_GPUS_PER_NODE} --node-ip-address ${WORKER_IP} --disable-usage-stats" &
-  done
-  wait
-fi
-
-RUNTIME_ENV_JSON=$(cat <<EOF_JSON
-{
-  "env_vars": {
-    "no_proxy": "localhost,127.0.0.1,0.0.0.0,${MASTER_ADDR}",
-    "GLOO_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "TP_SOCKET_IFNAME": "${SOCKET_IFNAME}",
-    "MASTER_ADDR": "${MASTER_ADDR}",
-    "PYTHONPATH": "/root/Megatron-LM/",
-    "CUDA_DEVICE_MAX_CONNECTIONS": "1",
-    "NCCL_NVLS_ENABLE": "${HAS_NVLINK}"
+ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
+
+# Build the runtime environment JSON with proper variable substitution
+RUNTIME_ENV_JSON="{
+  \"env_vars\": {
+    \"PYTHONPATH\": \"/root/Megatron-LM/\",
+    \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
+    \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
+    \"no_proxy\": \"${no_proxy}\"
   }
-}
-EOF_JSON
-)
+}"
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
    -- python3 train.py \
-   --actor-num-nodes "${ACTOR_NUM_NODES}" \
-   --actor-num-gpus-per-node "${ACTOR_NUM_GPUS_PER_NODE}" \
+   --actor-num-nodes 1 \
+   --actor-num-gpus-per-node 8 \
    --colocate \
-   "${MODEL_ARGS[@]}" \
-   "${CKPT_ARGS[@]}" \
-   "${ROLLOUT_ARGS[@]}" \
-   "${OPTIMIZER_ARGS[@]}" \
-   "${GRPO_ARGS[@]}" \
-   "${WANDB_ARGS[@]}" \
-   "${PERF_ARGS[@]}" \
-   "${EVAL_ARGS[@]}" \
-   "${SGLANG_ARGS[@]}" \
-   "${MISC_ARGS[@]}"
+   ${MODEL_ARGS[@]} \
+   ${CKPT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
+   ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
+   ${WANDB_ARGS[@]} \
+   ${PERF_ARGS[@]} \
+   ${EVAL_ARGS[@]} \
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]}
diff --git a/scripts/run-qwen3.5-35B-A3B-sft.sh b/scripts/run-qwen3.5-35B-A3B.sh
old mode 100644
new mode 100755
similarity index 54%
rename from scripts/run-qwen3.5-35B-A3B-sft.sh
rename to scripts/run-qwen3.5-35B-A3B.sh
index 6893133924..6cabc17e86
--- a/scripts/run-qwen3.5-35B-A3B-sft.sh
+++ b/scripts/run-qwen3.5-35B-A3B.sh
@@ -12,21 +12,12 @@ pkill -9 python
 
 set -ex
 
-# if base folder not set raise error
-if [ -z "${BASE_FOLDER}" ]; then
-  echo "BASE_FOLDER is not set. Please set it to the base directory of your checkpoints."
-  exit 1
-fi
-
-if [ -z "${MASTER_ADDR}" ]; then
-  echo "MASTER_ADDR is not set. Please set it to the master node address."
-  exit 1
-fi
-# export MASTER_ADDR="127.0.0.1"
-
 # will prevent ray from buffering stdout/stderr
 export PYTHONBUFFERED=16
 
+# unset proxy to avoid issues
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
 NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l)
 if [ "$NVLINK_COUNT" -gt 0 ]; then
     HAS_NVLINK=1
@@ -39,31 +30,40 @@ SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
 source "${SCRIPT_DIR}/models/qwen3.5-35B-A3B.sh"
 
 CKPT_ARGS=(
-   --hf-checkpoint ${BASE_FOLDER}/Qwen3.5-35B-A3B
-   --ref-load ${BASE_FOLDER}/Qwen3.5-35B-A3B_torch_dist
-   --load ${BASE_FOLDER}/Qwen3.5-35B-A3B_slime/
-   --save ${BASE_FOLDER}/Qwen3.5-35B-A3B_slime/
+   --hf-checkpoint /root/Qwen3.5-35B-A3B
+   --ref-load /root/Qwen3.5-35B-A3B_torch_dist
+   --load /root/Qwen3.5-35B-A3B_slime/
+   --save /root/Qwen3.5-35B-A3B_slime/
    --save-interval 20
 )
 
-SFT_ARGS=(
-   --rollout-function-path slime.rollout.sft_rollout.generate_rollout
-   --prompt-data ${BASE_FOLDER}/openhermes2_5.parquet
-   --input-key messages
+ROLLOUT_ARGS=(
+   --prompt-data /root/dapo-math-17k/dapo-math-17k.jsonl
+   --input-key prompt
+   --label-key label
+   --apply-chat-template
    --rollout-shuffle
-   --num-epoch 3
-   --rollout-batch-size 128
-   --global-batch-size 128
-
-   --loss-type sft_loss
-   --loss-mask-type qwen3_5
-   --calculate-per-token-loss
-   --disable-compute-advantages-and-returns
-   --debug-train-only
+   --rm-type deepscaler
+   --num-rollout 3000
+   --rollout-batch-size 32
+   --n-samples-per-prompt 8
+   --rollout-max-response-len 8192
+   --rollout-temperature 1
+
+   --global-batch-size 256
+   --balance-data
+)
+
+EVAL_ARGS=(
+   --eval-interval 20
+   --eval-prompt-data aime /root/aime-2024/aime-2024.jsonl
+   --n-samples-per-eval-prompt 16
+   --eval-max-response-len 16384
+   --eval-top-p 1
 )
 
 PERF_ARGS=(
-   --tensor-model-parallel-size 2
+   --tensor-model-parallel-size 4
    --sequence-parallel
    --pipeline-model-parallel-size 1
    --context-parallel-size 1
@@ -76,29 +76,53 @@ PERF_ARGS=(
 
    # --micro-batch-size 1
    --use-dynamic-batch-size
-   --max-tokens-per-gpu 8192
+   --max-tokens-per-gpu 20480
+)
+
+GRPO_ARGS=(
+   --advantage-estimator grpo
+   --use-kl-loss
+   --kl-loss-coef 0.00
+   --kl-loss-type low_var_kl
+   --entropy-coef 0.00
+   --eps-clip 0.2
+   --eps-clip-high 0.28
 )
 
 OPTIMIZER_ARGS=(
    --optimizer adam
-   --lr 1e-5
-   --lr-decay-style cosine
-   --min-lr 1e-6
-   --lr-warmup-fraction 0.1
+   --lr 1e-6
+   --lr-decay-style constant
    --weight-decay 0.1
    --adam-beta1 0.9
    --adam-beta2 0.98
-   
-   --use-distributed-optimizer
+
    --optimizer-cpu-offload
    --overlap-cpu-optimizer-d2h-h2d
    --use-precision-aware-optimizer
 )
 
 WANDB_ARGS=(
-   # --use-wandb
+   #--use-wandb
    # --wandb-project slime-dev
-   # --wandb-group qwen3.5-35B-sft
+   # --wandb-group qwen3.5-35B-A3B-test
+   # --wandb-key ${WANDB_KEY}
+)
+
+SGLANG_ARGS=(
+   --rollout-num-gpus-per-engine 8
+   --sglang-mem-fraction-static 0.7
+   --sglang-ep-size 8
+
+   --sglang-cuda-graph-bs 1 2 4 8 $(seq 16 8 256)
+
+   # mtp
+   --sglang-speculative-algorithm EAGLE
+   --sglang-speculative-num-steps 2
+   --sglang-speculative-eagle-topk 1
+   --sglang-speculative-num-draft-tokens 3
+
+   --sglang-max-running-requests 512
 )
 
 MISC_ARGS=(
@@ -112,28 +136,12 @@ MISC_ARGS=(
    --attention-backend flash
 
    --moe-token-dispatcher-type flex
-   --moe-enable-deepep
-)
-
-SPEC_ARGS=(
-#    --mtp-num-layers 1
-#    --enable-mtp-training
-#    --mtp-loss-scaling-factor 0.1
 )
 
 # launch the master node of ray in container
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
 export no_proxy="127.0.0.1,${MASTER_ADDR}"
 ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265
-for WORKER_IP in $(awk '{print $1}' /root/mpi_rack_hostfile); do
-  if [[ "$WORKER_IP" == "$MLP_WORKER_0_HOST" ]]; then
-    continue
-  fi
-  echo "Starting Ray worker on ${WORKER_IP}"
-  ssh root@"${WORKER_IP}" \
-    "pkill -9 sglang ; ray stop --force ; pkill -9 python ; ray start --address=${MASTER_ADDR}:6379 --num-gpus 8 --node-ip-address ${WORKER_IP} --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265" &
-done
-wait
-
 
 # Build the runtime environment JSON with proper variable substitution
 RUNTIME_ENV_JSON="{
@@ -141,23 +149,23 @@ RUNTIME_ENV_JSON="{
     \"PYTHONPATH\": \"/root/Megatron-LM/\",
     \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",
     \"NCCL_NVLS_ENABLE\": \"${HAS_NVLINK}\",
-    \"no_proxy\": \"${no_proxy}\",
-    \"MASTER_ADDR\": \"${MASTER_ADDR}\",
-    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\"
+    \"no_proxy\": \"${no_proxy}\"
   }
 }"
 
 ray job submit --address="http://127.0.0.1:8265" \
    --runtime-env-json="${RUNTIME_ENV_JSON}" \
-   -- python3 train_async.py \
+   -- python3 train.py \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 8 \
+   --colocate \
    ${MODEL_ARGS[@]} \
    ${CKPT_ARGS[@]} \
-   ${SFT_ARGS[@]} \
+   ${ROLLOUT_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
+   ${GRPO_ARGS[@]} \
    ${WANDB_ARGS[@]} \
    ${PERF_ARGS[@]} \
    ${EVAL_ARGS[@]} \
-   ${MISC_ARGS[@]} \
-   ${SPEC_ARGS[@]}
+   ${SGLANG_ARGS[@]} \
+   ${MISC_ARGS[@]}
diff --git a/setup.py b/setup.py
index c424a3a33d..91df81c5e4 100644
--- a/setup.py
+++ b/setup.py
@@ -32,11 +32,15 @@ def get_tag(self):
 setup(
     author="slime Team",
     name="slime",
-    version="0.2.4",
+    version="0.2.2",
     packages=find_packages(include=["slime*", "slime_plugins*"]),
     include_package_data=True,
     install_requires=_fetch_requirements("requirements.txt"),
-    extras_require={},
+    extras_require={
+        "fsdp": [
+            "torch>=2.0",
+        ]
+    },
     python_requires=">=3.10",
     classifiers=[
         "Programming Language :: Python :: 3.10",
diff --git a/setup_for_vllm.md b/setup_for_vllm.md
deleted file mode 100644
index ccdc8fd593..0000000000
--- a/setup_for_vllm.md
+++ /dev/null
@@ -1,21 +0,0 @@
-```
-docker pull slimerl/slime:latest
-```
-
-```
-docker run -itd --gpus all --ipc=host --shm-size=128g --net=host --privileged=true  --restart=always \
---ulimit memlock=-1 --ulimit stack=67108864 \
---ulimit nofile=65536:65536 \
---name DNAME  \
--it slimerl/slime:latest /bin/bash \
-
-```
-docker exec -it --user root DNAME bash
-```
-
-```
-pip install vllm=0.16 
-
-# for compatibility
-pip install numpy==1.26.4
-```
diff --git a/slime/backends/fsdp_utils/__init__.py b/slime/backends/fsdp_utils/__init__.py
new file mode 100644
index 0000000000..ba17c4725f
--- /dev/null
+++ b/slime/backends/fsdp_utils/__init__.py
@@ -0,0 +1,8 @@
+import logging
+
+from .actor import FSDPTrainRayActor
+from .arguments import fsdp_parse_args
+
+__all__ = ["fsdp_parse_args", "FSDPTrainRayActor"]
+
+logging.getLogger().setLevel(logging.WARNING)
diff --git a/slime/backends/fsdp_utils/actor.py b/slime/backends/fsdp_utils/actor.py
new file mode 100644
index 0000000000..5d9ff32a2f
--- /dev/null
+++ b/slime/backends/fsdp_utils/actor.py
@@ -0,0 +1,993 @@
+import logging
+import os
+import random
+from argparse import Namespace
+from itertools import accumulate
+
+import ray
+import torch
+import torch.distributed as dist
+from tqdm import tqdm
+from transformers import AutoConfig
+
+from slime.ray.train_actor import TrainRayActor
+from slime.utils import logging_utils, train_dump_utils, train_metric_utils
+from slime.utils.data import get_minimum_num_micro_batch_size, process_rollout_data
+from slime.utils.distributed_utils import get_gloo_group
+from slime.utils.logging_utils import init_tracking
+from slime.utils.memory_utils import clear_memory, print_memory
+from slime.utils.metric_utils import compute_rollout_step
+from slime.utils.misc import Box
+from slime.utils.ppo_utils import compute_approx_kl, compute_gspo_kl, compute_opsm_mask, compute_policy_loss
+from slime.utils.processing_utils import load_processor, load_tokenizer
+from slime.utils.profile_utils import TrainProfiler
+from slime.utils.timer import Timer, inverse_timer, timer, with_defer
+
+from . import checkpoint
+from .data_packing import pack_sequences, unpack_sequences
+from .lr_scheduler import get_lr_scheduler
+from .update_weight_utils import UpdateWeightFromDistributed, UpdateWeightFromTensor
+
+logger = logging.getLogger(__name__)
+
+
+class FSDPTrainRayActor(TrainRayActor):
+    """Simplified TrainRayActor for pure HF+FSDP training.
+
+    Responsibilities:
+      * Initialize model/tokenizer on rank0 sequentially to avoid race on cache
+      * Wrap model with FSDP
+      * Provide minimal train / save / update_weights hooks compatible with existing RayTrainGroup
+
+    Weight update strategy:
+      * Rank0 gathers state_dict (full) and broadcasts tensor-by-tensor.
+      * For small models this is fine; for larger models consider sharded state_dict type.
+    """
+
+    @with_defer(lambda: Timer().start("train_wait"))
+    def init(self, args: Namespace, role: str, with_ref: bool = False, with_opd_teacher: bool = False) -> int:  # type: ignore[override]
+        if with_opd_teacher:
+            raise NotImplementedError(
+                "On-policy distillation (OPD) with Megatron teacher is not supported in FSDP backend. "
+                "Please use the Megatron backend for OPD, or use --opd-type=sglang with an external teacher server."
+            )
+        super().init(args, role, with_ref, with_opd_teacher)
+
+        # Setup device mesh for data parallelism
+        self._setup_device_mesh()
+        torch.manual_seed(args.seed)
+
+        self.train_parallel_config = {
+            "dp_size": self.dp_size,
+        }
+
+        if self.args.debug_rollout_only:
+            return 0
+
+        self.fsdp_cpu_offload = getattr(self.args, "fsdp_cpu_offload", False)
+        # Offload train and fsdp cpu offload cannot be used together, fsdp_cpu_offload is more aggressive
+        if self.args.offload_train and self.fsdp_cpu_offload:
+            self.args.offload_train = False
+
+        self._enable_true_on_policy_optimizations(args)
+        if dist.get_rank() == 0:
+            init_tracking(args, primary=False)
+
+        if getattr(self.args, "start_rollout_id", None) is None:
+            self.args.start_rollout_id = 0
+
+        self.prof = TrainProfiler(args)
+
+        for i in range(dist.get_world_size()):
+            if i == dist.get_rank():
+                self.hf_config = AutoConfig.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
+                self.tokenizer = load_tokenizer(self.args.hf_checkpoint, trust_remote_code=True)
+                # Vision models have `vision_config` in the config
+                if hasattr(self.hf_config, "vision_config"):
+                    self.processor = load_processor(self.args.hf_checkpoint, trust_remote_code=True)
+            dist.barrier(group=get_gloo_group())
+
+        init_context = self._get_init_weight_context_manager()
+
+        with init_context():
+            model = self.get_model_cls().from_pretrained(
+                self.args.hf_checkpoint,
+                trust_remote_code=True,
+                attn_implementation=self.args.attn_implementation,
+            )
+
+        model.train()
+
+        full_state = model.state_dict()
+
+        model = apply_fsdp2(model, mesh=self.dp_mesh, cpu_offload=self.fsdp_cpu_offload, args=self.args)
+
+        model = self._fsdp2_load_full_state_dict(
+            model, full_state, self.dp_mesh, cpu_offload=True if self.fsdp_cpu_offload else None
+        )
+
+        self.model = model
+
+        if args.gradient_checkpointing:
+            self.model.gradient_checkpointing_enable()
+
+        if args.optimizer == "adam":
+            self.optimizer = torch.optim.AdamW(
+                self.model.parameters(),
+                lr=args.lr,
+                betas=(args.adam_beta1, args.adam_beta2),
+                eps=args.adam_eps,
+                weight_decay=args.weight_decay,
+            )
+        else:
+            raise ValueError(f"Unsupported optimizer: {args.optimizer}. Supported options: 'adam'")
+
+        # Initialize LR scheduler
+        self.lr_scheduler = get_lr_scheduler(args, self.optimizer)
+
+        self.global_step = 0
+        self.micro_step = 0
+
+        checkpoint_payload = checkpoint.load(self)
+
+        # Create separate ref model if needed (kept in CPU until needed)
+        self.ref_model = None
+        if with_ref:
+            self.ref_model = self._create_ref_model(args.ref_load)
+
+        self.weight_updater = (
+            UpdateWeightFromTensor(self.args, self.model)
+            if self.args.colocate
+            else UpdateWeightFromDistributed(self.args, self.model)
+        )
+
+        checkpoint.finalize_load(self, checkpoint_payload)
+
+        # Initialize data packing parameters
+        self.max_tokens_per_gpu = args.max_tokens_per_gpu  # From main arguments
+
+        if self.args.offload_train:
+            self.sleep()
+
+        self.prof.on_init_end()
+
+        return int(getattr(self.args, "start_rollout_id", 0))
+
+    def get_model_cls(self):
+        # Vision models have `vision_config` in the config
+        if hasattr(self.hf_config, "vision_config"):
+            from transformers import AutoModelForImageTextToText
+
+            return AutoModelForImageTextToText
+        else:
+            from transformers import AutoModelForCausalLM
+
+            return AutoModelForCausalLM
+
+    def _enable_true_on_policy_optimizations(self, args):
+        if args.true_on_policy_mode:
+            from sglang.srt.batch_invariant_ops import enable_batch_invariant_mode
+
+            from .models.qwen3_moe import apply_true_on_policy_patch_for_qwen3_moe
+
+            logger.info("FSDPTrainRayActor call enable_batch_invariant_mode for true-on-policy")
+            enable_batch_invariant_mode(
+                # In Qwen3, rope `inv_freq_expanded.float() @ position_ids_expanded.float()` uses bmm
+                # and disabling it will make it aligned
+                enable_bmm=False,
+            )
+
+            apply_true_on_policy_patch_for_qwen3_moe()
+        else:
+            from .models.qwen3_moe_hf import apply_fsdp_moe_patch
+
+            apply_fsdp_moe_patch()
+
+    def _setup_device_mesh(self) -> None:
+        """Setup device mesh for data parallelism."""
+        from torch.distributed.device_mesh import init_device_mesh
+
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # Pure data parallelism
+        self.dp_size = world_size
+        self.dp_rank = rank
+
+        # Create 1D device mesh for data parallelism
+        self.mesh = init_device_mesh("cuda", mesh_shape=(self.dp_size,), mesh_dim_names=("dp",))
+        self.dp_group = self.mesh.get_group("dp")
+        self.dp_mesh = self.mesh
+
+        logger.info(f"[Rank {rank}] Device mesh (1D): world_size={world_size}, dp_size={self.dp_size}")
+
+    def _get_init_weight_context_manager(self):
+        """Get context manager for model initialization.
+
+        Returns a callable that creates a context manager.
+        Uses meta device (no memory allocation) for non-rank-0 processes,
+        UNLESS tie_word_embeddings=True (which causes hangs with meta tensors).
+
+        Ref: verl/utils/fsdp_utils.py::get_init_weight_context_manager
+        NOTE: tie_word_embedding causes meta_tensor init to hang
+        """
+        from accelerate import init_empty_weights
+
+        # Check if model uses tied word embeddings (which doesn't work with meta tensors)
+        use_meta_tensor = not self.hf_config.tie_word_embeddings
+
+        def cpu_init_weights():
+            return torch.device("cpu")
+
+        if use_meta_tensor:
+            # Rank 0: CPU, others: meta device (memory efficient for large models)
+            return init_empty_weights if dist.get_rank() != 0 else cpu_init_weights
+        else:
+            logger.info(f"[Rank {dist.get_rank()}] tie_word_embeddings=True, loading full model to CPU on all ranks")
+            return cpu_init_weights
+
+    def _fsdp2_load_full_state_dict(self, model, full_state, device_mesh, cpu_offload):
+        """Load full state dict into FSDP2 model with efficient broadcast from rank 0.
+
+        This function loads weights from rank 0 and broadcasts to all other ranks,
+        avoiding the need for each rank to load the full model from disk.
+
+        Args:
+            model: FSDP2-wrapped model
+            full_state: State dict (only rank 0 has real weights, others have empty dict)
+            device_mesh: Device mesh for FSDP
+            cpu_offload: If not None, enables StateDictOptions cpu_offload
+
+        Ref:verl/utils/fsdp_utils.py::fsdp2_load_full_state_dict
+        """
+        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict
+
+        # Rank 0: move with weights, others: allocate empty tensors on device
+        if dist.get_rank() == 0:
+            model = model.to(device=torch.cuda.current_device(), non_blocking=True)
+        else:
+            # to_empty creates tensors on device without initializing memory
+            model = model.to_empty(device=torch.cuda.current_device())
+
+        is_cpu_offload = cpu_offload is not None
+        options = StateDictOptions(full_state_dict=True, cpu_offload=is_cpu_offload, broadcast_from_rank0=True)
+
+        set_model_state_dict(model, full_state, options=options)
+
+        # set_model_state_dict will not broadcast buffers, so we need to broadcast them manually.
+        for _name, buf in model.named_buffers():
+            dist.broadcast(buf, src=0)
+
+        if is_cpu_offload:
+            model.to("cpu", non_blocking=True)
+            for buf in model.buffers():
+                buf.data = buf.data.to(torch.cuda.current_device())
+
+        return model
+
+    @timer
+    def sleep(self) -> None:
+        """Pause CUDA memory for all tracked tensors."""
+        if not self.args.offload_train:
+            return
+
+        print_memory("before offload model")
+
+        self.model.cpu()
+        move_torch_optimizer(self.optimizer, "cpu")
+        clear_memory()
+        dist.barrier(group=get_gloo_group())
+        print_memory("after offload model")
+
+    @timer
+    def wake_up(self) -> None:
+        """Resume CUDA memory for all tracked tensors."""
+        if not self.args.offload_train:
+            return
+
+        self.model.cuda()
+        move_torch_optimizer(self.optimizer, "cuda")
+        dist.barrier(group=get_gloo_group())
+        print_memory("after wake_up model")
+
+    def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
+        """Delegate checkpoint saving to the shared checkpoint utilities."""
+        if self.args.debug_rollout_only or self.args.save is None:
+            return
+
+        assert not self.args.async_save, "FSDPTrainRayActor does not support async_save yet."
+        checkpoint.save(self, rollout_id)
+
+    def _compute_log_prob(
+        self,
+        model_tag: str,
+        packed_batches: list[dict[str, torch.Tensor]],
+        store_prefix: str = "",
+    ) -> dict[str, list[torch.Tensor]]:
+        """Compute token log-probabilities for a list of packed batches.
+
+        Parameters:
+            model_tag: Which parameters to use, e.g. "actor" or "ref".
+            packed_batches: A list of packed batch dictionaries produced by
+                `pack_sequences`, each containing at least `tokens` and
+                `position_ids`; may also include multimodal keys like `pixel_values`.
+            store_prefix: Prefix to use for keys in outputs (e.g., "ref_").
+
+        Returns:
+            A lightweight dictionary keyed by f"{store_prefix}log_probs". The
+            actual per-sequence results are written in-place into each element of
+            `packed_batches` under the same key and can be read back by callers.
+
+        Note:
+            Uses separate ref model when model_tag == "ref". The ref model is
+            loaded from CPU to GPU on-demand and offloaded back after use.
+        """
+        # Select which model to use
+        if model_tag == "ref" and self.ref_model is not None:
+            if not self.fsdp_cpu_offload:
+                self.model.cpu()
+                torch.cuda.empty_cache()
+                dist.barrier(group=get_gloo_group())
+
+            active_model = self.ref_model
+            active_model.eval()
+        else:
+            active_model = self.model
+
+        try:
+            rollout_data = {f"{store_prefix}log_probs": []}
+            with timer(f"{store_prefix}log_probs"), torch.no_grad():
+                for batch in self.prof.iterate_train_log_probs(
+                    tqdm(packed_batches, desc=f"{store_prefix}log_probs", disable=dist.get_rank() != 0)
+                ):
+                    model_args = self._get_model_inputs_args(batch)
+                    logits = active_model(**model_args).logits.squeeze(0).float()
+                    log_probs_result, entropy_result = get_logprob_and_entropy(
+                        logits=logits,
+                        target_tokens=batch["tokens"],
+                        allow_compile=not self.args.true_on_policy_mode,
+                        temperature=self.args.rollout_temperature,
+                    )
+                    batch[f"{store_prefix}log_probs"] = log_probs_result
+                    if store_prefix == "":
+                        batch["entropy"] = entropy_result
+            return rollout_data
+
+        finally:
+            # Restore actor model if it was offloaded
+            if model_tag == "ref" and self.ref_model is not None:
+                torch.cuda.empty_cache()
+                dist.barrier(group=get_gloo_group())
+
+                if not self.fsdp_cpu_offload:
+                    self.model.cuda()
+                    dist.barrier(group=get_gloo_group())
+
+    def _packed_data(
+        self, rollout_data: dict[str, list[torch.Tensor]]
+    ) -> tuple[list[dict[str, torch.Tensor]], list[int]]:
+        """Pack variable-length sequences for efficient processing.
+
+        Parameters:
+            rollout_data: Dictionary of lists containing sequence-level tensors
+                such as `tokens`, `loss_masks`, `rewards`, `response_lengths`,
+                `advantages`, `returns`, and optional `rollout_log_probs`.
+
+        Returns:
+            A pair `(packed_batches, grad_accum)` where `packed_batches` is a list
+            of packed batch dictionaries and `grad_accum` lists the micro-batch
+            indices at which to perform optimizer steps.
+        """
+        # Pack sequences efficiently
+        tokens = rollout_data["tokens"]
+
+        packed_batches = []
+        mbs_size_list = []
+        local_batch_size = self.args.global_batch_size // self.dp_size
+        assert (
+            self.args.global_batch_size % self.dp_size == 0
+        ), f"global_batch_size {self.args.global_batch_size} is not divisible by dp_world_size {self.dp_size}"
+        # Use global_batch_size for splitting when max_tokens_per_gpu is enabled
+        if self.args.use_dynamic_batch_size:
+            max_tokens = self.args.max_tokens_per_gpu
+
+            for i in range(0, len(tokens), local_batch_size):
+                mbs_size_list.append(
+                    get_minimum_num_micro_batch_size(
+                        [len(t) for t in rollout_data["tokens"][i : i + local_batch_size]],
+                        max_tokens,
+                    )
+                )
+            num_microbatches = torch.tensor(mbs_size_list, dtype=torch.int, device=torch.cuda.current_device())
+            dist.all_reduce(num_microbatches, op=dist.ReduceOp.MAX, group=self.dp_group)
+            num_microbatches = num_microbatches.tolist()
+        else:
+            num_microbatches = [self.args.global_batch_size // (self.args.micro_batch_size * self.dp_size)] * (
+                len(tokens) // local_batch_size
+            )
+
+        start = 0
+        for mbs_size in num_microbatches:
+            end = start + local_batch_size
+            packed_batches.extend(
+                pack_sequences(
+                    rollout_data["tokens"][start:end],
+                    rollout_data["loss_masks"][start:end],
+                    rollout_data["rewards"][start:end],
+                    rollout_data["raw_reward"][start:end],
+                    rollout_data["response_lengths"][start:end],
+                    rollout_data["advantages"][start:end],
+                    rollout_data["returns"][start:end],
+                    rollout_log_probs=(
+                        rollout_data["rollout_log_probs"][start:end] if "rollout_log_probs" in rollout_data else None
+                    ),
+                    multimodal_train_inputs=(
+                        rollout_data["multimodal_train_inputs"][start:end]
+                        if "multimodal_train_inputs" in rollout_data
+                        else None
+                    ),
+                    num_packs=mbs_size,
+                )
+            )
+            start = end
+        grad_accum = list(accumulate(num_microbatches))
+
+        return packed_batches, grad_accum
+
+    def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
+        """Run one training update over a rollout batch.
+
+        Parameters:
+            rollout_id: Monotonic id for logging.
+            rollout_data_ref: A Box handle wrapping a Ray object reference to a
+                dictionary with rollout tensors and metadata (e.g., `tokens`,
+                `loss_masks`, `rewards`, `response_lengths`, optional
+                `rollout_log_probs`, etc.). It will be fetched and partitioned
+                by `process_rollout_data` based on data-parallel rank/size.
+        """
+        if self.args.offload_train:
+            self.wake_up()
+
+        with inverse_timer("train_wait"), timer("train"):
+            rollout_data = process_rollout_data(self.args, rollout_data_ref, self.dp_rank, self.dp_size)
+            if self.args.debug_rollout_only:
+                return
+            self._train_core(rollout_id=rollout_id, rollout_data=rollout_data)
+
+        train_metric_utils.log_perf_data_raw(
+            rollout_id=rollout_id,
+            args=self.args,
+            is_primary_rank=dist.get_rank() == 0,
+            compute_total_fwd_flops=None,
+        )
+
+    def _log_rollout_data(self, rollout_id: int, rollout_data, packed_batches):
+        log_dict = {}
+        if "raw_reward" in rollout_data and dist.get_rank() == 0:
+            raw_reward_list = rollout_data["raw_reward"]
+            if raw_reward_list:
+                log_dict["rollout/raw_reward"] = sum(raw_reward_list) / len(raw_reward_list)
+
+        for metric_key in ["log_probs", "rollout_log_probs", "ref_log_probs", "advantages", "returns"]:
+            if metric_key not in packed_batches[0]:
+                continue
+            val = torch.tensor([0.0], device=torch.cuda.current_device())
+            for _mbs_id, batches in enumerate(packed_batches):
+                unpacked_batches = unpack_sequences(batches)
+                for unpacked_batch in unpacked_batches:
+                    if isinstance(unpacked_batch[metric_key], torch.Tensor):
+                        loss_masks_tensor = unpacked_batch["loss_masks"].to(device=torch.cuda.current_device())
+                        metric_tensor = unpacked_batch[metric_key].to(device=torch.cuda.current_device())
+                        val += (metric_tensor * loss_masks_tensor).sum() / loss_masks_tensor.sum().clamp_min(1)
+                    else:
+                        val += unpacked_batch[metric_key]
+            dist.all_reduce(val, op=dist.ReduceOp.SUM, group=self.dp_group)
+            log_dict[f"rollout/{metric_key}"] = (
+                val / (self.args.n_samples_per_prompt * self.args.rollout_batch_size)
+            ).item()
+        if dist.get_rank() == 0:
+            logger.info(f"rollout {rollout_id}: {log_dict}")
+            log_dict["rollout/step"] = compute_rollout_step(self.args, rollout_id)
+            logging_utils.log(self.args, log_dict, step_key="rollout/step")
+
+        if self.args.ci_test and self.args.true_on_policy_mode:
+            assert log_dict["rollout/log_probs"] == log_dict["rollout/rollout_log_probs"], (
+                f"CI check failed: true_on_policy_mode is enabled, but log_probs "
+                f"({log_dict['rollout/log_probs']}) != rollout_log_probs "
+                f"({log_dict['rollout/rollout_log_probs']})"
+            )
+
+    def _train_core(self, rollout_id: int, rollout_data) -> None:
+        if self.args.advantage_estimator in ["grpo", "gspo"]:
+            rollout_data["advantages"] = rollout_data["returns"] = [
+                torch.tensor([rollout_data["rewards"][i]] * rollout_data["response_lengths"][i])
+                for i in range(len(rollout_data["rewards"]))
+            ]
+        else:
+            raise NotImplementedError(f"Unsupported advantage_estimator {self.args.advantage_estimator}")
+
+        packed_batches, grad_accum = self._packed_data(rollout_data)
+
+        assert (
+            len(grad_accum) > 0
+        ), f"Invalid grad_accum {grad_accum} for micro_batch_size {self.args.micro_batch_size} and global_batch_size {self.args.global_batch_size}"
+
+        if self.ref_model is not None:
+            self._compute_log_prob("ref", packed_batches, store_prefix="ref_")
+
+        self._compute_log_prob("actor", packed_batches)
+        self._log_rollout_data(rollout_id, rollout_data, packed_batches)
+
+        with timer("actor_train"):
+            reported_accum: dict[str, list[torch.Tensor]] = {}
+            self.optimizer.zero_grad(set_to_none=True)
+            for mbs_id, packed_batch in self.prof.iterate_train_actor(
+                enumerate(tqdm(packed_batches, desc="actor_train", disable=dist.get_rank() != 0))
+            ):
+                self._train_step(
+                    packed_batch=packed_batch,
+                    reported_accum=reported_accum,
+                    mbs_id=mbs_id,
+                    grad_accum=grad_accum,
+                )
+
+        self.prof.step(rollout_id=rollout_id)
+
+        train_dump_utils.save_debug_train_data(self.args, rollout_id=rollout_id, rollout_data=rollout_data)
+
+        # Update ref model if needed (copy actor weights to ref)
+        if (
+            self.args.ref_update_interval is not None
+            and (rollout_id + 1) % self.args.ref_update_interval == 0
+            and self.ref_model is not None
+        ):
+            if dist.get_rank() == 0:
+                logger.info(f"Updating ref model at rollout_id {rollout_id}")
+            # Copy actor model state to ref model
+            actor_state = self.model.state_dict()
+            self.ref_model.load_state_dict(actor_state)
+            self.ref_model.cpu()
+
+    def _train_step(self, packed_batch, reported_accum, mbs_id, grad_accum):
+        # Prepare model inputs
+        model_args = self._get_model_inputs_args(packed_batch)
+        logits = self.model(**model_args).logits.squeeze(0).float()
+
+        # Compute log probs and entropy
+        log_probs, entropy_result = get_logprob_and_entropy(
+            logits=logits,
+            target_tokens=packed_batch["tokens"],
+            allow_compile=not self.args.true_on_policy_mode,
+            temperature=self.args.rollout_temperature,
+        )
+        packed_batch["cur_log_probs"] = log_probs
+        packed_batch["entropy"] = entropy_result
+
+        unpacked_batches = unpack_sequences(packed_batch)
+
+        old_log_prob_key = "rollout_log_probs" if self.args.use_rollout_logprobs else "log_probs"
+        missing_old_log_probs = [
+            idx
+            for idx, batch in enumerate(unpacked_batches)
+            if old_log_prob_key not in batch or not isinstance(batch[old_log_prob_key], torch.Tensor)
+        ]
+        if missing_old_log_probs:
+            raise KeyError(
+                f"{old_log_prob_key} must be provided as torch.Tensor for all microbatches when "
+                f"use_rollout_logprobs is set to {self.args.use_rollout_logprobs}. Missing in batches: {missing_old_log_probs}"
+            )
+        old_log_probs = torch.cat([batch[old_log_prob_key] for batch in unpacked_batches], dim=0)
+        log_probs = torch.cat([batch["cur_log_probs"] for batch in unpacked_batches], dim=0)
+        advantages = torch.cat([batch["advantages"] for batch in unpacked_batches], dim=0)
+        loss_masks = [batch["loss_masks"].to(device=log_probs.device) for batch in unpacked_batches]
+        response_lengths = [batch["response_lengths"] for batch in unpacked_batches]
+
+        advantages = advantages.to(device=log_probs.device)
+        old_log_probs = old_log_probs.to(device=log_probs.device)
+        ppo_kl = old_log_probs - log_probs
+
+        if self.args.use_opsm:
+            opsm_mask, opsm_clipfrac = compute_opsm_mask(
+                args=self.args,
+                full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches],
+                advantages=[batch["advantages"] for batch in unpacked_batches],
+                loss_masks=loss_masks,
+            )
+
+        if self.args.advantage_estimator == "gspo":
+            ppo_kl = compute_gspo_kl(
+                full_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                full_old_log_probs=[batch[old_log_prob_key] for batch in unpacked_batches],
+                local_log_probs=[batch["cur_log_probs"] for batch in unpacked_batches],
+                loss_masks=loss_masks,
+            )
+
+        pg_loss, pg_clipfrac = compute_policy_loss(ppo_kl, advantages, self.args.eps_clip, self.args.eps_clip_high)
+
+        if self.args.use_opsm:
+            pg_loss = pg_loss * opsm_mask
+
+        def _has_rollout_log_probs(batch) -> bool:
+            rollout_tensor = batch.get("rollout_log_probs")
+            return isinstance(rollout_tensor, torch.Tensor) and rollout_tensor.numel() > 0
+
+        has_rollout_log_probs = all(_has_rollout_log_probs(batch) for batch in unpacked_batches)
+        rollout_log_probs = (
+            torch.cat([batch["rollout_log_probs"] for batch in unpacked_batches], dim=0)
+            if has_rollout_log_probs
+            else None
+        )
+
+        if self.args.calculate_per_token_loss:
+            pg_loss = sum_of_token(pg_loss, response_lengths, loss_masks)
+            pg_clipfrac = sum_of_token(pg_clipfrac, response_lengths, loss_masks)
+            ppo_kl = sum_of_token(ppo_kl.abs(), response_lengths, loss_masks)
+        else:
+            pg_loss = sum_of_sample_mean(pg_loss, response_lengths, loss_masks)
+            pg_clipfrac = sum_of_sample_mean(pg_clipfrac, response_lengths, loss_masks)
+            ppo_kl = sum_of_sample_mean(ppo_kl.abs(), response_lengths, loss_masks)
+
+        # Only compare rollout vs. train log probs when they originate from different stages.
+        train_rollout_logprob_abs_diff = None
+        if not self.args.use_rollout_logprobs and rollout_log_probs is not None:
+            train_rollout_logprob_abs_diff = (old_log_probs - rollout_log_probs).abs()
+            train_rollout_logprob_abs_diff = sum_of_sample_mean(
+                train_rollout_logprob_abs_diff, response_lengths, loss_masks
+            ).detach()
+
+        entropy = torch.cat([batch["entropy"] for batch in unpacked_batches], dim=0)
+        entropy_loss = sum_of_sample_mean(entropy, response_lengths, loss_masks)
+
+        loss = pg_loss - self.args.entropy_coef * entropy_loss
+
+        if self.args.use_kl_loss:
+            ref_log_probs = torch.cat([batch["ref_log_probs"] for batch in unpacked_batches], dim=0)
+            importance_ratio = None
+            if self.args.use_unbiased_kl:
+                importance_ratio = torch.exp(log_probs - old_log_probs)
+            kl = compute_approx_kl(
+                log_probs,
+                ref_log_probs,
+                kl_loss_type=self.args.kl_loss_type,
+                importance_ratio=importance_ratio,
+            )
+            kl_loss = sum_of_sample_mean(kl, response_lengths, loss_masks)
+
+            loss = loss + self.args.kl_loss_coef * kl_loss
+
+        reported = {
+            "loss": loss.detach(),
+            "pg_loss": pg_loss.detach(),
+            "pg_clipfrac": pg_clipfrac.detach(),
+            "ppo_kl": ppo_kl.detach(),
+            "entropy_loss": entropy_loss.detach(),
+        }
+
+        if train_rollout_logprob_abs_diff is not None:
+            reported["train_rollout_logprob_abs_diff"] = train_rollout_logprob_abs_diff
+
+        if self.args.use_kl_loss:
+            reported["kl_loss"] = kl_loss.detach()
+
+        if self.args.use_opsm:
+            reported["opsm_clipfrac"] = opsm_clipfrac
+
+        # Scale loss for gradient accumulation
+        loss = loss * self.dp_size / self.args.global_batch_size
+        loss.backward()
+
+        # Accumulate reported metrics (store tensors for later mean)
+        for k, v in reported.items():
+            reported_accum.setdefault(k, []).append(v)
+
+        if (mbs_id + 1) in grad_accum:
+            # TODO: check if the grad norm is global grad norm.
+            grad_norm = torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.clip_grad)
+            # the grad norm used to be of DTensor
+            grad_norm = float(grad_norm)
+
+            self.optimizer.step()
+            # Update learning rate
+            self.lr_scheduler.step()
+            self.optimizer.zero_grad(set_to_none=True)
+            # Aggregate logs
+            aggregated = {k: torch.stack(v).sum().item() for k, v in reported_accum.items()}
+            # TODO: change this, this is slow.
+            reduced_aggregated = [None] * self.dp_size
+            dist.all_gather_object(reduced_aggregated, aggregated, group=self.dp_group)
+            aggregated = {}
+            for k in reported_accum.keys():
+                aggregated[k] = sum([r[k] for r in reduced_aggregated]) / (self.args.global_batch_size)
+            reported_accum.clear()
+            if dist.get_rank() == 0:
+                log_dict = {
+                    f"train/{k}": (val.item() if torch.is_tensor(val) else val) for k, val in aggregated.items()
+                }
+                log_dict["train/grad_norm"] = grad_norm
+
+                # Log learning rate per parameter group; use scheduler's last computed LRs
+                lr_values = self.lr_scheduler.get_last_lr()
+                for gid, _group in enumerate(self.optimizer.param_groups):
+                    log_dict[f"train/lr-pg_{gid}"] = lr_values[gid]
+
+                kl_info = ""
+                if self.args.use_kl_loss and "kl_loss" in aggregated:
+                    kl_info = f", kl_loss: {aggregated['kl_loss']:.4f}, kl_penalty: {aggregated['kl_loss'] * self.args.kl_loss_coef:.4f}"
+                    logger.info(kl_info)
+                logger.info(f"step {self.global_step}: {log_dict}")
+
+                log_dict["train/step"] = self.global_step
+                logging_utils.log(self.args, log_dict, step_key="train/step")
+            self.global_step += 1
+
+    @timer
+    def update_weights(self) -> None:  # type: ignore[override]
+        """Synchronize actor weights to rollout engines.
+
+        Handles both colocated and distributed update modes. In offload mode,
+        wakes up parameters as needed to perform the update.
+        """
+        if self.args.debug_train_only or self.args.debug_rollout_only:
+            return
+
+        rollout_engines, rollout_engine_lock, num_new_engines, engine_gpu_counts, engine_gpu_offsets = ray.get(
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
+        )
+        if num_new_engines > 0:
+            self.weight_updater.connect_rollout_engines(
+                rollout_engines,
+                rollout_engine_lock,
+                engine_gpu_counts=engine_gpu_counts,
+                engine_gpu_offsets=engine_gpu_offsets,
+            )
+            dist.barrier(group=get_gloo_group())
+            if dist.get_rank() == 0:
+                ray.get(self.rollout_manager.clear_num_new_engines.remote())
+
+        self.weight_updater.update_weights()
+
+        if self.args.ci_test and len(rollout_engines) > 0:
+            engine = random.choice(rollout_engines)
+            engine_version = ray.get(engine.get_weight_version.remote())
+            if str(engine_version) != str(self.weight_updater.weight_version):
+                raise RuntimeError(
+                    f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                )
+
+        clear_memory()
+
+    def _create_ref_model(self, ref_load_path: str | None):
+        """Create and initialize a separate reference model with FSDP2 CPUOffloadPolicy.
+
+        Parameters:
+            ref_load_path: Path to a directory containing a HF checkpoint. If
+                None, a ValueError is raised.
+
+        Returns:
+            FSDP2-wrapped ref model with CPU offload enabled
+
+        Note:
+            Creates a separate FSDP2 model instance for the reference model.
+            ALWAYS uses CPUOffloadPolicy for the reference model to save memory,
+            regardless of the actor model's CPU offload setting.
+        """
+        if ref_load_path is None:
+            raise ValueError("ref_load_path must be provided when loading reference model")
+
+        if os.path.isdir(ref_load_path):
+            logger.info(f"[Rank {dist.get_rank()}] Creating separate ref model from {ref_load_path}")
+
+            init_context = self._get_init_weight_context_manager()
+
+            with init_context():
+                ref_model = self.get_model_cls().from_pretrained(
+                    ref_load_path,
+                    trust_remote_code=True,
+                    attn_implementation=self.args.attn_implementation,
+                )
+
+            full_state = ref_model.state_dict()
+
+            # Always use CPUOffloadPolicy for reference, let FSDP2 handle the offload. It is faster than model.cpu().
+            ref_model = apply_fsdp2(ref_model, mesh=self.dp_mesh, cpu_offload=True, args=self.args)
+            ref_model = self._fsdp2_load_full_state_dict(ref_model, full_state, self.dp_mesh, cpu_offload=True)
+
+            logger.info(f"[Rank {dist.get_rank()}] Reference model created with FSDP2 CPUOffloadPolicy")
+            return ref_model
+        else:
+            raise NotImplementedError(f"Loading from checkpoint file {ref_load_path} not yet implemented")
+
+    def _get_model_inputs_args(self, packed_sequence: dict) -> dict:
+        input_ids = packed_sequence["tokens"].unsqueeze(0)
+        position_ids = packed_sequence["position_ids"].unsqueeze(0)
+
+        model_args = {
+            "input_ids": input_ids,
+            "position_ids": position_ids,
+            "attention_mask": None,
+        }
+        if packed_sequence.get("multimodal_train_inputs"):
+            model_args.update(packed_sequence["multimodal_train_inputs"])
+        return model_args
+
+
+def selective_log_softmax_raw(logits: torch.Tensor, input_ids: torch.Tensor) -> torch.Tensor:
+    """Fused version of the common `log_softmax -> gather` operation.
+
+    The fused version of this operation avoids the (potentially large) memory overhead
+    of allocating a new tensor to store the full logprobs.
+
+    Parameters:
+        logits: Tensor of shape [..., V] containing model logits.
+        input_ids: Tensor of shape [...] of token indices whose log-probabilities are gathered.
+
+    Returns:
+        Tensor of shape [...] containing the log-probabilities corresponding to `input_ids`.
+    """
+    logprobs = logits.log_softmax(dim=-1)
+    return torch.gather(logprobs, dim=-1, index=input_ids.unsqueeze(-1)).squeeze(-1)
+
+
+selective_log_softmax_compiled = torch.compile(dynamic=True)(selective_log_softmax_raw)
+
+
+def gather_log_probs_packed(
+    shifted_logits: torch.Tensor,
+    input_ids: torch.Tensor,
+    allow_compile: bool,
+    cu_seqlens: torch.Tensor | float | None = None,
+    temperature: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Gather next-token log probabilities for packed sequences.
+
+    Parameters:
+        logits: Model logits of shape [B, T, V] or [T, V].
+        input_ids: Token ids of shape [B, T] or [T].
+        cu_seqlens: Optional cumulative sequence lengths (unused here). Present
+            for API compatibility with callers.
+
+    Returns:
+        A tensor of shape [T-1] (or [B, T-1]) with log-probabilities of targets.
+    """
+    # Handle batch dimension - logits should be [batch_size, seq_len, vocab_size]
+    if shifted_logits.dim() == 3:
+        # Remove batch dimension for packed sequences
+        shifted_logits = shifted_logits.squeeze(0)
+        input_ids = input_ids.squeeze(0)
+
+    if temperature is not None:
+        shifted_logits = shifted_logits.div(temperature)
+
+    targets = input_ids[1:].to(device=shifted_logits.device)
+
+    # Gather log probs for targets
+    selective_log_softmax = selective_log_softmax_compiled if allow_compile else selective_log_softmax_raw
+    return selective_log_softmax(shifted_logits, targets)
+
+
+def get_logprob_and_entropy(
+    logits: torch.Tensor,
+    target_tokens: torch.Tensor,
+    allow_compile: bool,
+    temperature: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Compute log probabilities and entropy.
+
+    Parameters:
+        logits: Model output logits with shape [seq_len, vocab_size]
+        target_tokens: Target tokens with shape [seq_len]
+        allow_compile: Whether to allow compilation
+        temperature: Temperature parameter (optional)
+
+    Returns:
+        log_probs: Log probabilities with shape [seq_len - 1]
+        entropy: Entropy with shape [seq_len - 1]
+    """
+    shifted_logits = logits[:-1, :]
+    log_probs = gather_log_probs_packed(
+        shifted_logits, target_tokens, allow_compile=allow_compile, temperature=temperature
+    )
+    log_probs_full = torch.log_softmax(shifted_logits, dim=-1)
+    probs = torch.softmax(shifted_logits, dim=-1)
+    entropy = -(probs * log_probs_full).sum(dim=-1)
+    return log_probs, entropy
+
+
+def sum_of_sample_mean(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor:
+    """Compute sum of per-sample means across variable-length responses.
+
+    Parameters:
+        x: Flat tensor containing concatenated per-token values across samples.
+        response_lengths: Lengths of each sample's response segment in `x`.
+        loss_masks: Per-sample masks aligned with `response_lengths`.
+
+    Returns:
+        A scalar tensor equal to the sum over samples of the mean value within
+        each sample's response segment.
+    """
+    return sum(
+        [
+            (x_i * loss_mask_i).sum() / torch.clamp_min(loss_mask_i.sum(), 1)
+            for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+        ]
+    )
+
+
+@torch.no_grad()
+def move_torch_optimizer(optimizer, device):
+    """ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py"""
+    if not optimizer.state:
+        return
+
+    for param_group in optimizer.param_groups:
+        for param in param_group["params"]:
+            state = optimizer.state[param]
+            for key, value in state.items():
+                if isinstance(value, torch.Tensor):
+                    state[key] = value.to(device, non_blocking=True)
+
+    torch.cuda.synchronize()
+
+
+def apply_fsdp2(model, mesh=None, cpu_offload=False, args=None):
+    """Apply FSDP v2 to the model.
+
+    Args:
+        model: The model to wrap with FSDP
+        mesh: Optional DeviceMesh for FSDP. If None, uses all ranks.
+        cpu_offload: If True, offload parameters, gradients, and optimizer states
+            to CPU. The optimizer step will run on CPU. (Default: False)
+        args: Arguments containing precision settings (fp16/bf16)
+
+    Ref: https://github.com/volcengine/verl/blob/main/verl/utils/fsdp_utils.py
+    """
+    from torch.distributed.fsdp import CPUOffloadPolicy, MixedPrecisionPolicy, fully_shard
+
+    offload_policy = CPUOffloadPolicy() if cpu_offload else None
+
+    layer_cls_to_wrap = model._no_split_modules
+    assert len(layer_cls_to_wrap) > 0 and layer_cls_to_wrap[0] is not None
+
+    modules = [
+        module
+        for name, module in model.named_modules()
+        if module.__class__.__name__ in layer_cls_to_wrap
+        or (isinstance(module, torch.nn.Embedding) and not model.config.tie_word_embeddings)
+    ]
+
+    # Determine precision policy based on args
+    param_dtype = torch.bfloat16  # Default to bf16 as before
+    reduce_dtype = torch.float32
+
+    if args.fp16:
+        param_dtype = torch.float16
+
+    logger.info(f"FSDP MixedPrecision Policy: param_dtype={param_dtype}, reduce_dtype={reduce_dtype}")
+
+    fsdp_kwargs = {
+        "mp_policy": MixedPrecisionPolicy(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+        ),
+        "offload_policy": offload_policy,
+        "mesh": mesh,
+    }
+
+    # Apply FSDP to each module (offload_policy=None is equivalent to not passing it)
+    for module in modules:
+        fully_shard(module, **fsdp_kwargs)
+
+    # Apply FSDP to the top-level model
+    fully_shard(model, **fsdp_kwargs)
+
+    return model
+
+
+def sum_of_token(x: torch.Tensor, response_lengths: list[int], loss_masks: list[torch.Tensor]) -> torch.Tensor:
+    return sum(
+        [
+            (x_i * loss_mask_i).sum()
+            for x_i, loss_mask_i in zip(x.split(response_lengths, dim=0), loss_masks, strict=False)
+        ]
+    )
diff --git a/slime/backends/fsdp_utils/arguments.py b/slime/backends/fsdp_utils/arguments.py
new file mode 100644
index 0000000000..fcb83b9726
--- /dev/null
+++ b/slime/backends/fsdp_utils/arguments.py
@@ -0,0 +1,101 @@
+import argparse
+import dataclasses
+from dataclasses import dataclass
+
+import yaml
+
+
+@dataclass
+class FSDPArgs:
+    # Optim
+    optimizer: str = "adam"  # Optimizer type: "adam" (AdamW)
+    lr: float = 2e-5
+    lr_warmup_init: float = 0.0
+    min_lr: float = 0.0
+    lr_decay_style: str = "constant"
+    lr_decay_iters: int | None = None
+    lr_warmup_iters: int = 0
+    lr_warmup_fraction: float | None = None
+    lr_wsd_decay_iters: int | None = None
+    lr_wsd_decay_style: str | None = None
+    use_checkpoint_lr_scheduler: bool = True
+    override_lr_scheduler: bool = False
+    weight_decay: float = 0.0
+    adam_beta1: float = 0.9
+    adam_beta2: float = 0.95
+    adam_eps: float = 1e-8
+    warmup_ratio: float = 0.03
+
+    attn_implementation: str = "flash_attention_2"
+
+    # Logging
+    wandb_project: str = "slime-fsdp"
+    wandb_run_name: str | None = None
+
+    # Precision
+    gradient_checkpointing: bool = False
+    fp16: bool = False
+
+    # FSDP configuration
+    fsdp_state_dict_cpu_offload: bool = True  # If True, offload full state dict to CPU during collection.
+    fsdp_cpu_offload: bool = (
+        False  # If True, offload parameters, gradients, and optimizer states to CPU (optimizer runs on CPU)
+    )
+    fsdp_cpu_backend: str | None = (
+        "gloo"  # CPU backend for FSDP CPU offload (e.g., "gloo"). Set to None to disable hybrid backend.
+    )
+
+    deterministic_mode: bool = False  # This name must be the same as Megatron's
+
+    # Profile
+    record_memory_history: bool = False
+    memory_snapshot_path: str = "snapshot.pickle"
+    use_pytorch_profiler: bool = False
+    profile_step_start: int = 10
+    profile_step_end: int = 12
+    tensorboard_dir: str | None = None
+
+    # YAML bookkeeping
+    config: str | None = None
+
+
+def _parse_fsdp_cli(extra_args_provider=None, ignore_unknown_args=False):
+    parser = argparse.ArgumentParser("FSDP SFT Training (slime)")
+    parser.add_argument("--config", type=str, default=None, help="YAML config path")
+    for f in dataclasses.fields(FSDPArgs):
+        if f.name == "config":
+            continue
+
+        # Handle union types like int | None, str | None, etc.
+        if hasattr(f.type, "__args__"):  # Check if it's a Union type
+            # For T | None, use T as the type
+            non_none_types = [t for t in f.type.__args__ if t is not type(None)]
+            arg_type = non_none_types[0] if non_none_types else str
+        else:
+            arg_type = f.type
+
+        if arg_type is bool:
+            parser.add_argument(f"--{f.name.replace('_', '-')}", action="store_true")
+        else:
+            parser.add_argument(f"--{f.name.replace('_', '-')}", type=arg_type, default=f.default)
+
+    if extra_args_provider is not None:
+        parser = extra_args_provider(parser)
+    if ignore_unknown_args:
+        args, _ = parser.parse_known_args()
+    else:
+        args = parser.parse_args()
+    return args
+
+
+def fsdp_parse_args(extra_args_provider=None, ignore_unknown_args=False):
+    args = _parse_fsdp_cli(extra_args_provider, ignore_unknown_args=ignore_unknown_args)
+    if args.config:
+        with open(args.config) as f:
+            data = yaml.safe_load(f) or {}
+        for k, v in data.items():
+            if not hasattr(args, k):
+                setattr(args, k, v)
+    args.rank = 0  # Primary process rank for wandb initialization
+    args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node
+    return args
diff --git a/slime/backends/fsdp_utils/checkpoint.py b/slime/backends/fsdp_utils/checkpoint.py
new file mode 100644
index 0000000000..6daf7f982c
--- /dev/null
+++ b/slime/backends/fsdp_utils/checkpoint.py
@@ -0,0 +1,250 @@
+from __future__ import annotations
+
+import json
+import logging
+import time
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.distributed as dist
+import torch.distributed.checkpoint as dcp
+from torch.distributed.checkpoint.state_dict import get_state_dict, set_state_dict
+from torch.distributed.checkpoint.stateful import Stateful
+
+logger = logging.getLogger(__name__)
+
+
+class ModelState(Stateful):
+    """Wrapper for model state only."""
+
+    def __init__(self, model):
+        self.model = model
+
+    def state_dict(self):
+        model_state_dict, _ = get_state_dict(self.model, optimizers=[])
+        return {"model": model_state_dict}
+
+    def load_state_dict(self, state_dict):
+        set_state_dict(self.model, optimizers=[], model_state_dict=state_dict["model"], optim_state_dict=None)
+
+
+class OptimizerState(Stateful):
+    """Wrapper for optimizer state only."""
+
+    def __init__(self, model, optimizer):
+        self.model = model
+        self.optimizer = optimizer
+
+    def state_dict(self):
+        _, optimizer_state_dict = get_state_dict(self.model, optimizers=self.optimizer)
+        return {"optim": optimizer_state_dict}
+
+    def load_state_dict(self, state_dict):
+        set_state_dict(
+            self.model, optimizers=self.optimizer, model_state_dict=None, optim_state_dict=state_dict["optim"]
+        )
+
+
+class LRSchedulerState(Stateful):
+    """Wrapper for LR scheduler state only."""
+
+    def __init__(self, lr_scheduler):
+        self.lr_scheduler = lr_scheduler
+
+    def state_dict(self):
+        return {"lr_scheduler": self.lr_scheduler.state_dict()}
+
+    def load_state_dict(self, state_dict):
+        self.lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
+
+
+def _read_checkpoint_metadata(path: Path) -> dict[str, Any]:
+    if not path.exists():
+        return {}
+    try:
+        return json.loads(path.read_text())
+    except json.JSONDecodeError:
+        logger.warning(f"Failed to parse checkpoint metadata at {path}")
+        return {}
+
+
+def _write_checkpoint_metadata(path: Path, metadata: dict[str, Any]) -> None:
+    tmp_path = path.with_suffix(path.suffix + ".tmp")
+    tmp_path.write_text(json.dumps(metadata, indent=2, sort_keys=True))
+    tmp_path.replace(path)
+
+
+def load(actor: Any) -> dict[str, Any] | None:
+    """Load checkpoint from disk.
+
+    Loads model weights and optionally optimizer state from separate directories.
+    This allows loading weights without optimizer or deleting optimizer before loading.
+    """
+    load_root = getattr(actor.args, "load", None)
+    if load_root is None:
+        return None
+
+    root_path = Path(load_root).expanduser()
+    if not root_path.exists():
+        logger.info(f"[FSDP] Checkpoint directory {root_path} not found; skipping load.")
+        return None
+
+    target_step = getattr(actor.args, "ckpt_step", None)
+    if target_step is None:
+        tracker_file = root_path / "latest_checkpointed_iteration.txt"
+        if not tracker_file.exists():
+            logger.info(f"[FSDP] No tracker file at {tracker_file}; skipping load.")
+            return None
+        tracker_text = tracker_file.read_text().strip()
+        target_step = int(tracker_text)
+
+    checkpoint_dir = root_path / f"iter_{target_step:07d}"
+    model_dir = checkpoint_dir / "model"
+    optimizer_dir = checkpoint_dir / "optimizer"
+    lr_scheduler_dir = checkpoint_dir / "lr_scheduler"
+
+    if not model_dir.exists():
+        logger.info(f"[FSDP] Model checkpoint {model_dir} not found; skipping load.")
+        return None
+
+    # Load model weights (always)
+    model_state = ModelState(actor.model)
+    state_dict = {"model_state": model_state}
+
+    try:
+        dcp.load(state_dict=state_dict, checkpoint_id=str(model_dir))
+        logger.info(f"[FSDP] Loaded model from {model_dir}")
+    except Exception as e:
+        logger.error(f"[FSDP] Failed to load model from {model_dir}: {e}")
+        return None
+
+    # Load optimizer state (optional)
+    load_optimizer = not getattr(actor.args, "no_load_optim", False) and hasattr(actor, "optimizer")
+    if load_optimizer and optimizer_dir.exists():
+        optimizer_state = OptimizerState(actor.model, actor.optimizer)
+        optim_state_dict = {"optim_state": optimizer_state}
+        try:
+            dcp.load(state_dict=optim_state_dict, checkpoint_id=str(optimizer_dir))
+            logger.info(f"[FSDP] Loaded optimizer from {optimizer_dir}")
+        except Exception as e:
+            logger.warning(f"[FSDP] Failed to load optimizer from {optimizer_dir}: {e}")
+    elif load_optimizer:
+        logger.info(f"[FSDP] Optimizer checkpoint not found at {optimizer_dir}, skipping optimizer load.")
+
+    # Load LR scheduler state (optional)
+    load_lr_scheduler = hasattr(actor, "lr_scheduler") and lr_scheduler_dir.exists()
+    if load_lr_scheduler:
+        lr_scheduler_state = LRSchedulerState(actor.lr_scheduler)
+        lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state}
+        try:
+            dcp.load(state_dict=lr_scheduler_state_dict, checkpoint_id=str(lr_scheduler_dir))
+            logger.info(f"[FSDP] Loaded LR scheduler from {lr_scheduler_dir}")
+        except Exception as e:
+            logger.warning(f"[FSDP] Failed to load LR scheduler from {lr_scheduler_dir}: {e}")
+    elif hasattr(actor, "lr_scheduler"):
+        logger.info(f"[FSDP] LR scheduler checkpoint not found at {lr_scheduler_dir}, skipping LR scheduler load.")
+
+    rng_state = None
+    rng_path = checkpoint_dir / "rng.pt"
+    if rng_path.exists():
+        rng_state = torch.load(rng_path, map_location="cpu")
+
+    metadata = _read_checkpoint_metadata(checkpoint_dir / "meta.json")
+
+    return {
+        "rng": rng_state,
+        "metadata": metadata,
+        "iteration": target_step,
+    }
+
+
+def finalize_load(actor: Any, checkpoint_payload: dict[str, Any] | None) -> None:
+    if checkpoint_payload is None:
+        dist.barrier()
+        return
+
+    if checkpoint_payload.get("rng") is not None and not getattr(actor.args, "no_load_rng", False):
+        rng_state = checkpoint_payload["rng"]
+        if "torch" in rng_state:
+            torch.set_rng_state(rng_state["torch"])
+        if torch.cuda.is_available() and "cuda" in rng_state:
+            torch.cuda.set_rng_state_all(rng_state["cuda"])
+
+    metadata = checkpoint_payload.get("metadata") or {}
+    iteration = checkpoint_payload.get("iteration")
+    if metadata:
+        actor.global_step = int(metadata.get("global_step", actor.global_step))
+        actor.micro_step = int(metadata.get("micro_step", actor.micro_step))
+        next_rollout = metadata.get("next_rollout_id")
+        if next_rollout is not None:
+            actor.args.start_rollout_id = next_rollout
+    elif iteration is not None:
+        if getattr(actor.args, "start_rollout_id", None) is None:
+            actor.args.start_rollout_id = iteration
+
+    torch.cuda.synchronize()
+    dist.barrier()
+
+
+def save(actor: Any, iteration: int) -> None:
+    """Save checkpoint to disk.
+
+    Saves model weights and optimizer state to separate directories.
+    This allows loading weights without optimizer or deleting optimizer before loading.
+    """
+    torch.cuda.synchronize()
+
+    base_dir = Path(actor.args.save).expanduser()
+    step_id = iteration + 1
+    checkpoint_dir = base_dir / f"iter_{step_id:07d}"
+    model_dir = checkpoint_dir / "model"
+    optimizer_dir = checkpoint_dir / "optimizer"
+    lr_scheduler_dir = checkpoint_dir / "lr_scheduler"
+
+    if dist.get_rank() == 0:
+        checkpoint_dir.mkdir(parents=True, exist_ok=True)
+        model_dir.mkdir(parents=True, exist_ok=True)
+        optimizer_dir.mkdir(parents=True, exist_ok=True)
+        lr_scheduler_dir.mkdir(parents=True, exist_ok=True)
+    dist.barrier()
+
+    # Save model weights
+    model_state = ModelState(actor.model)
+    state_dict = {"model_state": model_state}
+    dcp.save(state_dict, checkpoint_id=str(model_dir))
+
+    # Save optimizer state (skip if --no-save-optim is set)
+    save_optimizer_state = not getattr(actor.args, "no_save_optim", False)
+    if save_optimizer_state and hasattr(actor, "optimizer") and actor.optimizer is not None:
+        optimizer_state = OptimizerState(actor.model, actor.optimizer)
+        optim_state_dict = {"optim_state": optimizer_state}
+        dcp.save(optim_state_dict, checkpoint_id=str(optimizer_dir))
+
+    # Save LR scheduler state (skip if --no-save-optim is set)
+    if save_optimizer_state and hasattr(actor, "lr_scheduler") and actor.lr_scheduler is not None:
+        lr_scheduler_state = LRSchedulerState(actor.lr_scheduler)
+        lr_scheduler_state_dict = {"lr_scheduler_state": lr_scheduler_state}
+        dcp.save(lr_scheduler_state_dict, checkpoint_id=str(lr_scheduler_dir))
+
+    if dist.get_rank() == 0:
+        rng_state = {"torch": torch.get_rng_state()}
+        rng_state["cuda"] = torch.cuda.get_rng_state_all()
+        torch.save(rng_state, checkpoint_dir / "rng.pt")
+
+        metadata = {
+            "iteration": step_id,
+            "rollout_id": iteration,
+            "next_rollout_id": iteration + 1,
+            "global_step": actor.global_step,
+            "micro_step": actor.micro_step,
+            "world_size": dist.get_world_size(),
+            "timestamp": time.time(),
+        }
+        _write_checkpoint_metadata(checkpoint_dir / "meta.json", metadata)
+
+        tracker_file = base_dir / "latest_checkpointed_iteration.txt"
+        tracker_file.write_text(str(step_id))
+        logger.info(f"[FSDP] Saved checkpoint to {checkpoint_dir}")
+
+    dist.barrier()
diff --git a/slime/backends/fsdp_utils/data_packing.py b/slime/backends/fsdp_utils/data_packing.py
new file mode 100644
index 0000000000..180e5d71b8
--- /dev/null
+++ b/slime/backends/fsdp_utils/data_packing.py
@@ -0,0 +1,196 @@
+"""Data packing utilities for FSDP backend to reduce padding overhead."""
+
+import math
+
+import torch
+
+from slime.utils.seqlen_balancing import get_seqlen_balanced_partitions
+
+
+def pack_sequences(
+    tokens: list[list[int]],
+    loss_masks: list[list[int]],
+    rewards: list[float],
+    raw_rewards: list,
+    response_lengths: list[int],
+    advantages: list[float],
+    returns: list[float],
+    rollout_log_probs: list[list[float]] | None = None,
+    multimodal_train_inputs: list[dict] | None = None,
+    max_tokens_per_gpu: int | None = None,
+    num_packs: int | None = None,
+) -> list[dict]:
+    """
+    Pack sequences into dense batches with cumulative sequence lengths.
+
+    Args:
+        tokens: List of token sequences
+        loss_masks: List of loss masks
+        rewards: List of rewards per sequence
+        raw_rewards: List of raw rewards per sequence
+        response_lengths: List of response lengths per sequence
+        advantages: List of advantages per sequence
+        returns: List of returns per sequence
+        rollout_log_probs: List of rollout log probabilities per sequence
+        multimodal_train_inputs: List of dict of multimodal tensors for training per sequence
+        max_tokens_per_gpu: Maximum tokens per GPU pack
+        num_packs: Explicit number of packs to create
+
+    Returns:
+        List of packed batches with tokens, masks, cu_seqlens, rewards, raw_rewards, response_lengths, advantages, returns
+    """
+    if not tokens:
+        return []
+
+    seq_lengths = [len(t) for t in tokens]
+
+    # Determine number of packs and use balanced partitioning
+    if num_packs:
+        k_partitions = num_packs
+    elif max_tokens_per_gpu:
+        total_tokens = sum(seq_lengths)
+        k_partitions = max(1, math.ceil(total_tokens / max_tokens_per_gpu))
+    else:
+        k_partitions = 1
+
+    # Use balanced partitioning for optimal load distribution
+    partitions = get_seqlen_balanced_partitions(
+        seq_lengths, k_partitions=k_partitions, equal_size=False  # Allow variable sizes for better balance
+    )
+
+    # Pack each partition
+    result = []
+    for indices in partitions:
+        # Build cumulative sequence lengths
+        cu_seqlens = [0]
+        flat_tokens = []
+        flat_masks = []
+        flat_positionids = []
+        flat_advantages = []
+        flat_returns = []
+        flat_rollout_log_probs = []
+
+        for i in indices:
+            seq_tokens = tokens[i]
+            seq_mask = loss_masks[i]
+            seq_positionids = list(range(len(seq_tokens)))
+
+            flat_tokens.extend(seq_tokens)
+            flat_positionids.extend(seq_positionids)
+            flat_masks.extend(seq_mask)
+            flat_advantages.extend(advantages[i])
+            flat_returns.extend(returns[i])
+            if rollout_log_probs:
+                flat_rollout_log_probs.extend(rollout_log_probs[i])
+            cu_seqlens.append(cu_seqlens[-1] + len(seq_tokens))
+
+        packed_batch = {
+            "tokens": torch.tensor(flat_tokens, dtype=torch.long),
+            "loss_masks": torch.tensor(flat_masks, dtype=torch.int),
+            "position_ids": torch.tensor(flat_positionids, dtype=torch.int),
+            "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32),
+            "rewards": torch.tensor([rewards[i] for i in indices], dtype=torch.float32),
+            "raw_reward": [raw_rewards[i] for i in indices],
+            "response_lengths": [response_lengths[i] for i in indices],
+            "advantages": torch.tensor(flat_advantages, dtype=torch.float32),
+            "returns": torch.tensor(flat_returns, dtype=torch.float32),
+            "rollout_log_probs": torch.tensor(
+                flat_rollout_log_probs, dtype=torch.float32, device=torch.cuda.current_device()
+            ),
+        }
+
+        # Collect and add multimodal training tensors for this partition
+        if multimodal_train_inputs:
+            multimodal_data = {}  # key -> concatenated tensor
+            multimodal_num_items = {}  # key -> list of item counts per sequence
+            for i in indices:
+                if multimodal_train_inputs[i] is None:
+                    continue
+                for key, mm_tensor in multimodal_train_inputs[i].items():
+                    if key not in multimodal_data:
+                        multimodal_data[key] = mm_tensor
+                        multimodal_num_items[key] = [mm_tensor.size(0)]
+                    else:
+                        multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
+                        multimodal_num_items[key].append(mm_tensor.size(0))
+            packed_batch["multimodal_train_inputs"] = multimodal_data
+            packed_batch["multimodal_num_items"] = multimodal_num_items
+
+        result.append(packed_batch)
+
+    return result
+
+
+def unpack_sequences(packed_batch: dict) -> list[dict]:
+    """
+    Unpack sequences from a packed batch.
+
+    Args:
+        packed_batch: Packed batch
+
+    Returns:
+        List of unpacked batches
+    """
+
+    cu_seqlens = packed_batch["cu_seqlens"]
+    num_sequences = len(cu_seqlens) - 1
+    response_lengths = packed_batch["response_lengths"]
+    multimodal_num_items = packed_batch.get("multimodal_num_items", {})
+
+    instances = []
+
+    # Calculate pad_length by counting trailing zeros
+    tokens = packed_batch["tokens"]
+    nonzero_indices = (tokens != 0).nonzero(as_tuple=True)[0]
+    if len(nonzero_indices) > 0:
+        # Last non-zero index, pad_length is everything after it
+        pad_length = len(tokens) - nonzero_indices[-1].item() - 1
+    else:
+        pad_length = 0  # No padding if no non-zero tokens (or all zeros)
+    for i in range(num_sequences):
+        start_idx = cu_seqlens[i].item()
+        end_idx = cu_seqlens[i + 1].item()
+        instance = {}
+
+        # Copy any additional attributes that might exist in the packed batch
+        for key, value in packed_batch.items():
+            if key not in instance:
+                # Skip multimodal_num_items - it's metadata
+                if key == "multimodal_num_items":
+                    continue
+                # Handle multimodal_train_inputs dict: split each tensor using multimodal_num_items
+                elif key == "multimodal_train_inputs" and isinstance(value, dict):
+                    instance[key] = {}
+                    for mm_key, mm_tensor in value.items():
+                        if mm_key in multimodal_num_items:
+                            num_items_list = multimodal_num_items[mm_key]
+                            start_mm_idx = sum(num_items_list[:i])
+                            end_mm_idx = start_mm_idx + num_items_list[i]
+                            if num_items_list[i] > 0:
+                                instance[key][mm_key] = mm_tensor[start_mm_idx:end_mm_idx]
+                # For tensor attributes, we need to slice them appropriately
+                elif isinstance(value, torch.Tensor):
+                    if key in ["log_probs", "ref_log_probs", "cur_log_probs", "entropy"]:
+                        # These are computed from logits[:-1] so they have length seq_len-1
+                        instance[key] = value[
+                            end_idx - 1 - response_lengths[i] - pad_length : end_idx - 1 - pad_length
+                        ]
+                    elif key == "rollout_log_probs":
+                        # rollout_log_probs is packed based on response_lengths, so slice differently
+                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
+                    elif key in ["tokens", "position_ids"]:
+                        # For other tensor attributes, try to slice them
+                        if len(value) > start_idx:
+                            instance[key] = value[start_idx:end_idx]
+                        else:
+                            raise ValueError(f"Attribute {key} is not found in the packed batch")
+                    elif key in ["loss_masks", "advantages", "returns"]:
+                        instance[key] = value[sum(response_lengths[:i]) : sum(response_lengths[: i + 1])]
+                elif isinstance(value, list):
+                    instance[key] = value[i]
+                else:
+                    raise ValueError(f"Attribute {key} is not found in the packed batch")
+
+        instances.append(instance)
+
+    return instances
diff --git a/slime/backends/fsdp_utils/kernels/__init__.py b/slime/backends/fsdp_utils/kernels/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/slime/backends/fsdp_utils/kernels/fused_experts.py b/slime/backends/fsdp_utils/kernels/fused_experts.py
new file mode 100644
index 0000000000..d1c02aae8a
--- /dev/null
+++ b/slime/backends/fsdp_utils/kernels/fused_experts.py
@@ -0,0 +1,381 @@
+from __future__ import annotations
+
+import torch
+import triton.language as tl
+from sglang.srt.layers.moe.fused_moe_triton.fused_moe import (
+    invoke_fused_moe_kernel,
+    moe_align_block_size,
+    moe_sum_reduce,
+    silu_and_mul,
+)
+
+from .fused_moe_triton_backward_kernels import invoke_fused_moe_backward_kernel
+
+
+class GateUpProjFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        num_tokens, _ = hidden_states.shape
+        E, N, _ = w1.shape
+        # We execute the fused_moe kernel in chunks to circumvent this issue:
+        # https://github.com/vllm-project/vllm/issues/5938
+        CHUNK_SIZE = 64 * 1024
+
+        # default deterministic config
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+        }
+
+        topk = topk_ids.shape[1]
+
+        intermediate_cache1 = torch.empty(
+            (num_tokens * topk, N),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+            curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+            cur_intermediate_cache1 = intermediate_cache1[begin_chunk_idx * topk : end_chunk_idx * topk]
+
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            invoke_fused_moe_kernel(
+                curr_hidden_states,
+                w1,
+                None,
+                cur_intermediate_cache1,
+                None,
+                None,
+                None,
+                curr_topk_weights,
+                curr_topk_ids,
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                False,
+                topk_ids.shape[1],
+                config,
+                compute_type=tl.bfloat16,
+                use_fp8_w8a8=False,
+                use_int8_w8a8=False,
+                use_int8_w8a16=False,
+                use_int4_w4a16=False,
+                per_channel_quant=False,
+                block_shape=None,
+                c_sorted=False,
+                filter_expert=True,
+            )
+
+        ctx.save_for_backward(hidden_states, w1, topk_weights, topk_ids)
+        ctx.config = config
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache1
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for GateUpProjFunction using Triton kernels.
+
+        Args:
+            grad_output: shape (num_tokens * topk, N)
+
+        Returns:
+            (grad_hidden_states, grad_w1, grad_topk_weights, None)
+        """
+
+        hidden_states, w1, topk_weights, topk_ids = ctx.saved_tensors
+        config = ctx.config
+        num_tokens = ctx.num_tokens
+        topk = ctx.topk
+
+        E, N, D_in = w1.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_hidden_states = torch.zeros_like(hidden_states)
+        grad_w1 = torch.zeros_like(w1)
+        # GateUpProj stage doesn't need topk_weights gradient
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx * topk : end_chunk_idx * topk]
+
+            # Get aligned metadata
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            # Prepare gradient buffer for this chunk
+            curr_grad_hidden_states = torch.zeros_like(curr_hidden_states)
+            curr_grad_w1 = torch.zeros_like(w1)
+
+            # Call Triton backward kernel with MUL_ROUTED_WEIGHT=False
+            # Use chunk of hidden_states to match sorted_token_ids indices
+            invoke_fused_moe_backward_kernel(
+                grad_output=curr_grad_output,
+                input=curr_hidden_states,  # Use chunk of hidden_states to match sorted_token_ids
+                weight=w1,
+                grad_input=curr_grad_hidden_states,
+                grad_weight=curr_grad_w1,
+                grad_topk_weights=None,  # Not needed for GateUpProj
+                topk_weights=curr_topk_weights,
+                topk_ids=curr_topk_ids,
+                sorted_token_ids=sorted_token_ids,
+                expert_ids=expert_ids,
+                num_tokens_post_padded=num_tokens_post_padded,
+                mul_routed_weight=False,
+                top_k=topk,
+                config=config,
+                compute_type=tl.bfloat16,
+            )
+
+            # Accumulate gradients
+            grad_hidden_states[begin_chunk_idx:end_chunk_idx] += curr_grad_hidden_states
+            grad_w1 += curr_grad_w1
+
+        return grad_hidden_states, grad_w1, grad_topk_weights, None
+
+
+class SiluAndMulFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, intermediate_cache1: torch.Tensor):
+        num_tokens, N = intermediate_cache1.shape
+        intermediate_cache2 = torch.empty(
+            (num_tokens, N // 2),
+            device=intermediate_cache1.device,
+            dtype=intermediate_cache1.dtype,
+        )
+        silu_and_mul(intermediate_cache1.view(-1, N), intermediate_cache2)
+
+        ctx.save_for_backward(intermediate_cache1)
+        return intermediate_cache2
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (intermediate_cache1,) = ctx.saved_tensors
+        N = intermediate_cache1.shape[-1]
+        x1, x2 = intermediate_cache1.view(-1, N).chunk(2, dim=-1)
+        silu_x1 = torch.nn.functional.silu(x1)
+
+        sig = torch.sigmoid(x1)
+        dsilu_dx1 = sig + x1 * sig * (1 - sig)
+        grad_x1 = grad_output * x2 * dsilu_dx1
+        grad_x2 = grad_output * silu_x1
+        grad_input = torch.cat([grad_x1, grad_x2], dim=-1)
+
+        return grad_input.view_as(intermediate_cache1)
+
+
+class DownProjFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        intermediate_cache2: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        num_tokens, _ = intermediate_cache2.shape
+        topk = topk_ids.shape[1]
+        num_tokens //= topk
+        E, _, _ = w2.shape
+        # We execute the fused_moe kernel in chunks to circumvent this issue:
+        # https://github.com/vllm-project/vllm/issues/5938
+        CHUNK_SIZE = 64 * 1024
+
+        # default deterministic config
+        config = {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 64,
+            "BLOCK_SIZE_K": 32,
+            "GROUP_SIZE_M": 8,
+        }
+
+        intermediate_cache3 = torch.empty(
+            (num_tokens, topk, w2.shape[1]),
+            device=intermediate_cache2.device,
+            dtype=intermediate_cache2.dtype,
+        )
+
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+            cur_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
+            cur_intermediate_cache3 = intermediate_cache3[begin_chunk_idx:end_chunk_idx]
+
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+            invoke_fused_moe_kernel(
+                cur_intermediate_cache2,
+                w2,
+                None,
+                cur_intermediate_cache3,
+                None,
+                None,
+                None,
+                curr_topk_weights,
+                curr_topk_ids,
+                sorted_token_ids,
+                expert_ids,
+                num_tokens_post_padded,
+                True,
+                1,
+                config,
+                compute_type=tl.bfloat16,
+                use_fp8_w8a8=False,
+                use_int8_w8a8=False,
+                use_int8_w8a16=False,
+                use_int4_w4a16=False,
+                per_channel_quant=False,
+                block_shape=None,
+                a_use_tma=False,
+                b_use_tma=False,
+            )
+
+        ctx.save_for_backward(intermediate_cache2, w2, topk_weights, topk_ids)
+        ctx.config = config
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache3
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for DownProjFunction using Triton kernels.
+
+        Args:
+            grad_output: shape (num_tokens, topk, hidden_size)
+
+        Returns:
+            (grad_intermediate_cache2, grad_w2, grad_topk_weights, None)
+        """
+        intermediate_cache2, w2, topk_weights, topk_ids = ctx.saved_tensors
+        config = ctx.config
+        num_tokens = ctx.num_tokens
+        topk = ctx.topk
+
+        E, hidden_size, intermediate_size = w2.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_intermediate_cache2 = torch.zeros_like(intermediate_cache2)
+        grad_w2 = torch.zeros_like(w2)
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx:end_chunk_idx]
+
+            # Get aligned metadata
+            sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
+                curr_topk_ids, config["BLOCK_SIZE_M"], E
+            )
+
+            # Prepare gradient buffers for this chunk
+            curr_grad_intermediate_cache2 = torch.zeros_like(curr_intermediate_cache2)
+            curr_grad_w2 = torch.zeros_like(w2)
+            curr_grad_topk_weights = torch.zeros_like(curr_topk_weights)
+
+            # Call Triton backward kernel with MUL_ROUTED_WEIGHT=True
+            # Note: Use top_k=1 to match forward pass indexing
+            invoke_fused_moe_backward_kernel(
+                grad_output=curr_grad_output,
+                input=curr_intermediate_cache2,
+                weight=w2,
+                grad_input=curr_grad_intermediate_cache2,
+                grad_weight=curr_grad_w2,
+                grad_topk_weights=curr_grad_topk_weights,
+                topk_weights=curr_topk_weights,
+                topk_ids=curr_topk_ids,
+                sorted_token_ids=sorted_token_ids,
+                expert_ids=expert_ids,
+                num_tokens_post_padded=num_tokens_post_padded,
+                mul_routed_weight=True,
+                top_k=1,
+                config=config,
+                compute_type=tl.bfloat16,
+            )
+
+            # Accumulate gradients
+            grad_intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk] = curr_grad_intermediate_cache2
+            grad_w2 += curr_grad_w2
+            grad_topk_weights[begin_chunk_idx:end_chunk_idx] = curr_grad_topk_weights
+
+        return grad_intermediate_cache2, grad_w2, grad_topk_weights, None
+
+
+class MoeSumReduceFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        intermediate_cache3: torch.Tensor,
+        hidden_states_shape,
+    ):
+        out_hidden_states = torch.empty(
+            hidden_states_shape, device=intermediate_cache3.device, dtype=intermediate_cache3.dtype
+        )
+        moe_sum_reduce(
+            intermediate_cache3,
+            out_hidden_states,
+            1.0,
+        )
+        ctx.save_for_backward(intermediate_cache3)
+        return out_hidden_states
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (intermediate_cache3,) = ctx.saved_tensors
+        return grad_output.unsqueeze(1).expand_as(intermediate_cache3), None
diff --git a/slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py b/slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py
new file mode 100644
index 0000000000..3333478491
--- /dev/null
+++ b/slime/backends/fsdp_utils/kernels/fused_moe_triton_backward_kernels.py
@@ -0,0 +1,540 @@
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def fused_moe_backward_input_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    weight_ptr,
+    grad_input_ptr,
+    grad_topk_weights_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_we,
+    stride_wn,
+    stride_wk,
+    stride_gim,
+    stride_gik,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_input.
+
+    Forward: output = input @ weight.T (optionally multiplied by topk_weights)
+    Backward: grad_input = grad_output @ weight (optionally multiplied by topk_weights)
+
+    This kernel computes: grad_input[token] = sum_over_N(grad_output[token, n] * weight[expert, n, :])
+    If MUL_ROUTED_WEIGHT: grad_input[token] *= topk_weights[token]
+
+    Parallelization: Similar to forward, parallel over M and N dimensions, loop over K.
+    """
+    # Map program ids to blocks (parallel over M and N, similar to forward)
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid_m * BLOCK_SIZE_M < num_tokens_post_padded:
+        # Load token information
+        offs_token_id = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+        offs_token = tl.load(sorted_token_ids_ptr + offs_token_id)
+        offs_token = offs_token.to(tl.int64)
+        token_mask = offs_token < num_valid_tokens
+
+        # Get expert ID for this block
+        off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+        # Only process if expert is valid
+        if off_experts != -1:
+            # Initialize offsets for N dimension (current block)
+            offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+            offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+            # Load grad_output block: shape (BLOCK_SIZE_M, BLOCK_SIZE_N)
+            grad_output_ptrs = grad_output_ptr + (offs_token[:, None] * stride_gom + offs_n[None, :] * stride_gon)
+            grad_out = tl.load(
+                grad_output_ptrs,
+                mask=token_mask[:, None] & (offs_n[None, :] < N),
+                other=0.0,
+            )
+
+            # Apply topk_weights to grad_output if needed
+            if MUL_ROUTED_WEIGHT:
+                moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
+                grad_out = grad_out * moe_weight[:, None]
+
+            # Iterate over K dimension
+            for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+                # Current K offsets
+                curr_offs_k = k * BLOCK_SIZE_K + offs_k
+
+                # Load weight block: shape (BLOCK_SIZE_N, BLOCK_SIZE_K)
+                # weight: shape (E, N, K)
+                weight_ptrs = (
+                    weight_ptr
+                    + off_experts * stride_we
+                    + offs_n[:, None] * stride_wn
+                    + curr_offs_k[None, :] * stride_wk
+                )
+                w = tl.load(
+                    weight_ptrs,
+                    mask=(offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
+                    other=0.0,
+                )
+
+                # Compute contribution: grad_out @ weight
+                # grad_out: (BLOCK_SIZE_M, BLOCK_SIZE_N)
+                # w: (BLOCK_SIZE_N, BLOCK_SIZE_K)
+                # result: (BLOCK_SIZE_M, BLOCK_SIZE_K)
+                contribution = tl.dot(grad_out, w)
+
+                # Atomic add to grad_input because different N blocks contribute to same K
+                grad_input_ptrs = grad_input_ptr + (
+                    (offs_token[:, None] // top_k) * stride_gim + curr_offs_k[None, :] * stride_gik
+                )
+                grad_input_mask = token_mask[:, None] & (curr_offs_k[None, :] < K)
+                tl.atomic_add(grad_input_ptrs, contribution.to(compute_type), mask=grad_input_mask)
+
+
+@triton.jit
+def fused_moe_backward_weight_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    input_ptr,
+    grad_weight_ptr,
+    topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_im,
+    stride_ik,
+    stride_gwe,
+    stride_gwn,
+    stride_gwk,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    MUL_ROUTED_WEIGHT: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_weight.
+
+    Forward: output = input @ weight.T (optionally multiplied by topk_weights)
+    Backward: grad_weight = input.T @ grad_output (optionally multiplied by topk_weights)
+
+    This kernel computes: grad_weight[expert, n, k] = sum_over_tokens(input[token, k] * grad_output[token, n])
+    If MUL_ROUTED_WEIGHT: the accumulation is weighted by topk_weights[token]
+
+    Parallelization: Parallel over M and N dimensions with grouping, loop over K.
+    """
+    # Map program ids to blocks (parallel over M and N with grouping, similar to forward and backward_input)
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
+        return
+
+    # Get expert ID for this M block
+    expert_id = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
+
+    # Only process if expert is valid
+    if expert_id == -1:
+        return
+
+    # Load token information for this M block
+    offs_m = tl.arange(0, BLOCK_SIZE_M)
+    offs_token_id = pid_m * BLOCK_SIZE_M + offs_m.to(tl.int64)
+    offs_token = tl.load(
+        sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
+    )
+    offs_token = offs_token.to(tl.int64)
+    token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
+
+    # Clamp offs_token to valid range
+    offs_token_clamped = tl.where(token_mask, offs_token, 0)
+
+    # Determine input token indices based on MUL_ROUTED_WEIGHT
+    if MUL_ROUTED_WEIGHT:
+        input_token_idx = offs_token_clamped
+        input_mask = token_mask
+    else:
+        input_token_idx = offs_token_clamped // top_k
+        num_input_tokens = num_valid_tokens // top_k
+        input_mask = token_mask & (input_token_idx < num_input_tokens)
+
+    # Load topk_weights if needed
+    if MUL_ROUTED_WEIGHT:
+        moe_weight = tl.load(topk_weights_ptr + offs_token_clamped, mask=token_mask, other=0.0)
+
+    # Current N offset for this program
+    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)
+
+    # Load grad_output for this N block: shape (M, BLOCK_SIZE_N)
+    # grad_output is always indexed by sorted_token_ids (offs_token_clamped)
+    # because it has shape (num_tokens * topk, N)
+    grad_output_ptrs = grad_output_ptr + (offs_token_clamped[:, None] * stride_gom + offs_n[None, :] * stride_gon)
+    grad_out = tl.load(
+        grad_output_ptrs,
+        mask=token_mask[:, None] & (offs_n[None, :] < N),
+        other=0.0,
+    )
+
+    # Apply topk_weights if needed
+    if MUL_ROUTED_WEIGHT:
+        grad_out = grad_out * moe_weight[:, None]
+
+    # Zero out padding tokens
+    token_mask_col = token_mask[:, None]
+    grad_out = grad_out * token_mask_col
+
+    # Iterate over K blocks and accumulate
+    for k_block in range(tl.cdiv(K, BLOCK_SIZE_K)):
+        offs_k = k_block * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K).to(tl.int64)
+
+        # Load input for this K block
+        input_ptrs = input_ptr + (input_token_idx[:, None] * stride_im + offs_k[None, :] * stride_ik)
+        inp = tl.load(
+            input_ptrs,
+            mask=input_mask[:, None] & (offs_k[None, :] < K),
+            other=0.0,
+        )
+
+        # Zero out padding tokens - use input_mask for input, token_mask for grad_output
+        input_mask_col = input_mask[:, None]
+        inp = inp * input_mask_col
+
+        # Compute grad_weight contribution: grad_out.T @ inp
+        grad_w_contribution = tl.dot(grad_out.T, inp)
+
+        # Write back using atomic add
+        grad_weight_ptrs = (
+            grad_weight_ptr + expert_id * stride_gwe + offs_n[:, None] * stride_gwn + offs_k[None, :] * stride_gwk
+        )
+        grad_weight_mask = (offs_n[:, None] < N) & (offs_k[None, :] < K)
+        tl.atomic_add(grad_weight_ptrs, grad_w_contribution.to(compute_type), mask=grad_weight_mask)
+
+
+@triton.jit
+def fused_moe_backward_topk_weights_kernel(
+    # Pointers to matrices
+    grad_output_ptr,
+    input_ptr,
+    weight_ptr,
+    grad_topk_weights_ptr,
+    sorted_token_ids_ptr,
+    expert_ids_ptr,
+    num_tokens_post_padded_ptr,
+    # Matrix dimensions
+    N,
+    K,
+    EM,
+    num_valid_tokens,
+    # Strides
+    stride_gom,
+    stride_gon,
+    stride_im,
+    stride_ik,
+    stride_we,
+    stride_wn,
+    stride_wk,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    top_k: tl.constexpr,
+    compute_type: tl.constexpr,
+):
+    """
+    Backward kernel for computing grad_topk_weights.
+
+    Forward: output = topk_weights * (input @ weight.T)
+    Backward: grad_topk_weights = sum(grad_output * (input @ weight.T))
+
+    This kernel computes the gradient of topk_weights by computing the dot product
+    of grad_output with the forward output before weight multiplication.
+    """
+    # Map program id to token block
+    pid = tl.program_id(axis=0)
+
+    # Check bounds
+    num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr)
+
+    # Only process if this block is valid
+    if pid * BLOCK_SIZE_M < num_tokens_post_padded:
+        # Load token information
+        offs_token_id = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M).to(tl.int64)
+        offs_token = tl.load(
+            sorted_token_ids_ptr + offs_token_id, mask=offs_token_id < num_tokens_post_padded, other=num_valid_tokens
+        )
+        offs_token = offs_token.to(tl.int64)
+        token_mask = (offs_token_id < num_tokens_post_padded) & (offs_token < num_valid_tokens)
+
+        # Clamp offs_token to valid range for safe pointer arithmetic
+        offs_token_clamped = tl.where(token_mask, offs_token, 0)
+
+        # Get expert ID for this block
+        off_experts = tl.load(expert_ids_ptr + pid).to(tl.int64)
+
+        # Only process if expert is valid
+        if off_experts != -1:
+            # Initialize offsets
+            offs_n = tl.arange(0, BLOCK_SIZE_N)
+            offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+            # Accumulator for grad_topk_weights
+            accumulator = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)
+
+            # Iterate over N and K dimensions to compute forward output and gradient
+            for n in range(0, tl.cdiv(N, BLOCK_SIZE_N)):
+                # Current N offset
+                curr_offs_n = n * BLOCK_SIZE_N + offs_n
+
+                # Load grad_output block: (M, N)
+                grad_output_ptrs = grad_output_ptr + (
+                    offs_token_clamped[:, None] * stride_gom + curr_offs_n[None, :] * stride_gon
+                )
+                grad_out = tl.load(
+                    grad_output_ptrs,
+                    mask=token_mask[:, None] & (curr_offs_n[None, :] < N),
+                    other=0.0,
+                )
+
+                # Compute forward output for this N block: input @ weight[:, n, :].T
+                forward_output_n = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+
+                for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+                    # Current K offset
+                    curr_offs_k = k * BLOCK_SIZE_K + offs_k
+
+                    # Load input block: (M, K)
+                    input_ptrs = input_ptr + (
+                        (offs_token_clamped[:, None] // top_k) * stride_im + curr_offs_k[None, :] * stride_ik
+                    )
+                    inp = tl.load(
+                        input_ptrs,
+                        mask=token_mask[:, None] & (curr_offs_k[None, :] < K),
+                        other=0.0,
+                    )
+
+                    # Load weight block: (N, K)
+                    weight_ptrs = (
+                        weight_ptr
+                        + off_experts * stride_we
+                        + curr_offs_n[:, None] * stride_wn
+                        + curr_offs_k[None, :] * stride_wk
+                    )
+                    w = tl.load(
+                        weight_ptrs,
+                        mask=(curr_offs_n[:, None] < N) & (curr_offs_k[None, :] < K),
+                        other=0.0,
+                    )
+
+                    # Accumulate forward output: input @ weight.T
+                    # inp: (M, K), w.T: (K, N) -> (M, N)
+                    forward_output_n += tl.dot(inp, w.T)
+
+                # Compute contribution to grad_topk_weights: sum(grad_out * forward_output)
+                # Sum over N dimension
+                accumulator += tl.sum(grad_out * forward_output_n, axis=1)
+
+            # Write back grad_topk_weights using atomic add with clamped token indices
+            tl.atomic_add(grad_topk_weights_ptr + offs_token_clamped, accumulator.to(compute_type), mask=token_mask)
+
+
+def invoke_fused_moe_backward_kernel(
+    grad_output: torch.Tensor,
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    grad_input: torch.Tensor,
+    grad_weight: torch.Tensor,
+    grad_topk_weights: torch.Tensor | None,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    sorted_token_ids: torch.Tensor,
+    expert_ids: torch.Tensor,
+    num_tokens_post_padded: torch.Tensor,
+    mul_routed_weight: bool,
+    top_k: int,
+    config: dict[str, Any],
+    compute_type: tl.dtype,
+) -> None:
+    """
+    Invoke the fused MOE backward kernels to compute gradients.
+
+    Args:
+        grad_output: Gradient of output, shape (num_tokens * topk, N) or (num_tokens, topk, N)
+        input: Input tensor, shape (num_tokens, K)
+        weight: Weight tensor, shape (E, N, K)
+        grad_input: Output gradient for input, shape (num_tokens, K)
+        grad_weight: Output gradient for weight, shape (E, N, K)
+        grad_topk_weights: Output gradient for topk_weights, shape (num_tokens, topk) or None
+        topk_weights: Top-K routing weights, shape (num_tokens, topk)
+        topk_ids: Top-K expert IDs, shape (num_tokens, topk)
+        sorted_token_ids: Sorted token IDs
+        expert_ids: Expert IDs for each block
+        num_tokens_post_padded: Number of tokens after padding
+        mul_routed_weight: Whether to multiply by routing weights
+        top_k: Number of experts per token
+        config: Kernel configuration
+        compute_type: Computation data type
+    """
+    assert topk_weights.stride(1) == 1
+    assert sorted_token_ids.stride(0) == 1
+
+    # Flatten grad_output if needed
+    # Before: (num_tokens, topk, hidden_size)
+    # After: (num_tokens * topk, hidden_size)
+    if grad_output.ndim == 3:
+        grad_output = grad_output.reshape(-1, grad_output.shape[-1])
+
+    E, N, K = weight.shape
+
+    # ===================== Compute grad_input =====================
+    def grid_input(META):
+        return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
+
+    fused_moe_backward_input_kernel[grid_input](
+        grad_output,
+        weight,
+        grad_input,
+        grad_topk_weights if grad_topk_weights is not None else grad_input,  # dummy pointer
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        sorted_token_ids.shape[0],
+        grad_output.shape[0],
+        grad_output.stride(0),
+        grad_output.stride(1),
+        weight.stride(0),
+        weight.stride(1),
+        weight.stride(2),
+        grad_input.stride(0),
+        grad_input.stride(1),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        **config,
+    )
+
+    # ===================== Compute grad_weight =====================
+    # Initialize grad_weight to zero
+    grad_weight.zero_()
+
+    # Use same grid configuration as forward kernel: encode both M and N dimensions
+    def grid_weight(META):
+        return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
+
+    fused_moe_backward_weight_kernel[grid_weight](
+        grad_output,
+        input,
+        grad_weight,
+        topk_weights,
+        sorted_token_ids,
+        expert_ids,
+        num_tokens_post_padded,
+        N,
+        K,
+        sorted_token_ids.shape[0],
+        grad_output.shape[0],
+        grad_output.stride(0),
+        grad_output.stride(1),
+        input.stride(0),
+        input.stride(1),
+        grad_weight.stride(0),
+        grad_weight.stride(1),
+        grad_weight.stride(2),
+        MUL_ROUTED_WEIGHT=mul_routed_weight,
+        top_k=top_k,
+        compute_type=compute_type,
+        **config,
+    )
+
+    # ===================== Compute grad_topk_weights (if needed) =====================
+    if mul_routed_weight and grad_topk_weights is not None:
+
+        def grid_topk(META):
+            return (triton.cdiv(sorted_token_ids.shape[0], META["BLOCK_SIZE_M"]),)
+
+        fused_moe_backward_topk_weights_kernel[grid_topk](
+            grad_output,
+            input,
+            weight,
+            grad_topk_weights.view(-1),
+            sorted_token_ids,
+            expert_ids,
+            num_tokens_post_padded,
+            N,
+            K,
+            sorted_token_ids.shape[0],
+            grad_output.shape[0],
+            grad_output.stride(0),
+            grad_output.stride(1),
+            input.stride(0),
+            input.stride(1),
+            weight.stride(0),
+            weight.stride(1),
+            weight.stride(2),
+            top_k=top_k,
+            compute_type=compute_type,
+            BLOCK_SIZE_M=config["BLOCK_SIZE_M"],
+            BLOCK_SIZE_N=config["BLOCK_SIZE_N"],
+            BLOCK_SIZE_K=config["BLOCK_SIZE_K"],
+        )
diff --git a/slime/backends/fsdp_utils/lr_scheduler.py b/slime/backends/fsdp_utils/lr_scheduler.py
new file mode 100644
index 0000000000..9bbe1efd16
--- /dev/null
+++ b/slime/backends/fsdp_utils/lr_scheduler.py
@@ -0,0 +1,195 @@
+# Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved.
+
+"""Learning rate scheduler for FSDP training."""
+
+import logging
+import math
+
+import torch
+from torch.optim.lr_scheduler import LRScheduler
+from typing_extensions import override
+
+logger = logging.getLogger(__name__)
+
+
+class FSDPLRScheduler(LRScheduler):
+    """Learning rate scheduler for FSDP training.
+
+    Args:
+        optimizer (torch.optim.Optimizer): The optimizer to be used.
+        init_lr (float): Initial learning rate.
+        max_lr (float): Maximum learning rate.
+        min_lr (float): Minimum learning rate.
+        lr_warmup_steps (int): Number of warmup steps.
+        lr_decay_steps (int): Number of decay steps.
+        lr_decay_style (str): Decay style for learning rate.
+        use_checkpoint_lr_scheduler (bool, optional): Whether to use the checkpoint values
+            for the lr scheduler.
+        override_lr_scheduler (bool, optional): Whether to override the lr scheduler values
+            with the class values.
+        wsd_decay_steps (int, optional): Number of weight decay decay steps.
+        lr_wsd_decay_style (str, optional): Decay style for learning rate during weight decay decay
+            steps.
+        last_epoch (int, optional): The index of last epoch. Default: -1.
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        init_lr: float,
+        max_lr: float,
+        min_lr: float,
+        lr_warmup_steps: int,
+        lr_decay_steps: int,
+        lr_decay_style: str,
+        use_checkpoint_lr_scheduler: bool | None = True,
+        override_lr_scheduler: bool | None = False,
+        wsd_decay_steps: int | None = None,
+        lr_wsd_decay_style: str | None = None,
+        last_epoch: int = -1,
+    ) -> None:
+        # Store our custom parameters
+        self.init_lr = init_lr
+        self.max_lr = float(max_lr)
+        self.min_lr = min_lr
+        assert self.min_lr >= 0.0
+        assert self.max_lr >= self.min_lr
+        assert self.init_lr <= self.max_lr
+
+        self.lr_warmup_steps = lr_warmup_steps
+        self.lr_decay_steps = lr_decay_steps
+        self.wsd_decay_steps = wsd_decay_steps
+        self.lr_wsd_decay_style = lr_wsd_decay_style
+
+        assert self.lr_decay_steps > 0
+        assert self.lr_warmup_steps < self.lr_decay_steps
+
+        self.lr_decay_style = lr_decay_style
+        if self.lr_decay_style == "WSD":
+            assert self.wsd_decay_steps is not None
+
+        self.override_lr_scheduler = override_lr_scheduler
+        self.use_checkpoint_lr_scheduler = use_checkpoint_lr_scheduler
+
+        if self.override_lr_scheduler:
+            assert not self.use_checkpoint_lr_scheduler, "both override and use-checkpoint are set."
+
+        # Initialize parent class
+        super().__init__(optimizer, last_epoch)
+
+        logger.info(f"> learning rate decay style: {self.lr_decay_style}")
+
+    def _get_lr_for_group(self, param_group: dict) -> float:
+        """Compute learning rate for a specific parameter group.
+
+        Args:
+            param_group (dict): parameter group from the optimizer.
+
+        Returns:
+            float: learning rate for this parameter group.
+        """
+        max_lr = param_group.get("max_lr", self.max_lr)
+        min_lr = param_group.get("min_lr", self.min_lr)
+
+        # Use linear warmup for the initial part.
+        if self.lr_warmup_steps > 0 and self.last_epoch <= self.lr_warmup_steps:
+            return self.init_lr + ((max_lr - self.init_lr) * float(self.last_epoch) / float(self.lr_warmup_steps))
+
+        # If the learning rate is constant, just return the initial value.
+        if self.lr_decay_style == "constant":
+            return max_lr
+
+        # For any steps larger than `self.lr_decay_steps`, use `min_lr`.
+        if self.last_epoch > self.lr_decay_steps:
+            return min_lr
+
+        # If we are done with the warmup period, use the decay style.
+        if self.lr_decay_style == "inverse-square-root":
+            warmup_steps = max(self.lr_warmup_steps, 1)
+            num_steps = max(self.last_epoch, 1)
+            lr = max_lr * warmup_steps**0.5 / (num_steps**0.5)
+            return max(min_lr, lr)
+
+        num_steps_ = self.last_epoch - self.lr_warmup_steps
+        decay_steps_ = self.lr_decay_steps - self.lr_warmup_steps
+        decay_ratio = float(num_steps_) / float(decay_steps_)
+        assert decay_ratio >= 0.0
+        assert decay_ratio <= 1.0
+
+        delta_lr = max_lr - min_lr
+        coeff = None
+
+        if self.lr_decay_style == "linear":
+            coeff = 1.0 - decay_ratio
+        elif self.lr_decay_style == "cosine":
+            coeff = 0.5 * (math.cos(math.pi * decay_ratio) + 1.0)
+        elif self.lr_decay_style == "WSD":
+            wsd_anneal_start_ = self.lr_decay_steps - self.wsd_decay_steps
+            if self.last_epoch <= wsd_anneal_start_:
+                coeff = 1.0
+            else:
+                wsd_steps = self.last_epoch - wsd_anneal_start_
+                wsd_decay_ratio = float(wsd_steps) / float(self.wsd_decay_steps)
+                if self.lr_wsd_decay_style == "linear":
+                    coeff = 1.0 - wsd_decay_ratio
+                elif self.lr_wsd_decay_style == "cosine":
+                    coeff = 0.5 * (math.cos(math.pi * wsd_decay_ratio) + 1.0)
+                elif self.lr_wsd_decay_style == "exponential":
+                    coeff = (2.0 * math.pow(0.5, wsd_decay_ratio)) - 1.0
+                elif self.lr_wsd_decay_style == "minus_sqrt":
+                    coeff = 1.0 - math.sqrt(wsd_decay_ratio)
+        else:
+            raise Exception(f"{self.lr_decay_style} decay style is not supported.")
+
+        assert coeff is not None
+        return min_lr + coeff * delta_lr
+
+    @override
+    def get_lr(self) -> list[float]:
+        """Compute the learning rates for each parameter group.
+
+        Returns:
+            list[float]: A list of learning rates, one for each parameter group.
+        """
+        return [self._get_lr_for_group(group) for group in self.optimizer.param_groups]
+
+
+def get_lr_scheduler(args, optimizer: torch.optim.Optimizer) -> FSDPLRScheduler:
+    """Create and configure the learning-rate scheduler.
+
+    This configures iteration-based schedules derived from the global batch size
+    and run-time arguments.
+
+    Args:
+        args: Training/runtime arguments (namespace).
+        optimizer (torch.optim.Optimizer): Optimizer bound to the model.
+
+    Returns:
+        FSDPLRScheduler: Initialized scheduler bound to ``optimizer``.
+    """
+    args.train_iters = args.num_rollout * args.rollout_batch_size * args.n_samples_per_prompt // args.global_batch_size
+    if args.lr_decay_iters is None:
+        args.lr_decay_iters = args.train_iters
+    lr_decay_steps = args.lr_decay_iters
+    wsd_decay_steps = None
+    if args.lr_wsd_decay_iters is not None:
+        wsd_decay_steps = args.lr_wsd_decay_iters
+    if args.lr_warmup_fraction is not None:
+        lr_warmup_steps = args.lr_warmup_fraction * lr_decay_steps
+    else:
+        lr_warmup_steps = args.lr_warmup_iters
+    lr_scheduler = FSDPLRScheduler(
+        optimizer,
+        init_lr=args.lr_warmup_init,
+        max_lr=args.lr,
+        min_lr=args.min_lr,
+        lr_warmup_steps=lr_warmup_steps,
+        lr_decay_steps=lr_decay_steps,
+        lr_decay_style=args.lr_decay_style,
+        use_checkpoint_lr_scheduler=args.use_checkpoint_lr_scheduler,
+        override_lr_scheduler=args.override_lr_scheduler,
+        wsd_decay_steps=wsd_decay_steps,
+        lr_wsd_decay_style=args.lr_wsd_decay_style,
+    )
+
+    return lr_scheduler
diff --git a/slime/backends/fsdp_utils/models/__init__.py b/slime/backends/fsdp_utils/models/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/slime/backends/fsdp_utils/models/qwen3_moe.py b/slime/backends/fsdp_utils/models/qwen3_moe.py
new file mode 100644
index 0000000000..2933f62a39
--- /dev/null
+++ b/slime/backends/fsdp_utils/models/qwen3_moe.py
@@ -0,0 +1,125 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeMLP
+
+from slime.backends.fsdp_utils.kernels.fused_experts import (
+    DownProjFunction,
+    GateUpProjFunction,
+    MoeSumReduceFunction,
+    SiluAndMulFunction,
+)
+
+
+def fused_experts_impl(
+    hidden_states: torch.Tensor,
+    w1: torch.Tensor,
+    w2: torch.Tensor,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+):
+    assert hidden_states.shape[1] == w1.shape[2], "Hidden size mismatch"
+    assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
+    assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
+    assert w1.is_contiguous(), "Expert weights1 must be contiguous"
+    assert w2.is_contiguous(), "Expert weights2 must be contiguous"
+    assert hidden_states.dtype in [torch.bfloat16]
+
+    intermediate_cache1 = GateUpProjFunction.apply(
+        hidden_states,
+        w1,
+        topk_weights,
+        topk_ids,
+    )
+    intermediate_cache2 = SiluAndMulFunction.apply(intermediate_cache1)
+    intermediate_cache3 = DownProjFunction.apply(
+        intermediate_cache2,
+        w2,
+        topk_weights,
+        topk_ids,
+    )
+    output_hidden_states = MoeSumReduceFunction.apply(
+        intermediate_cache3,
+        hidden_states.shape,
+    )
+    return output_hidden_states
+
+
+class StandardDispatcher:
+    def __init__(self, num_experts: int, num_local_experts: int):
+        self.moe_ep_size = 1
+        self.num_experts = num_experts
+        self.num_local_experts = num_local_experts
+        self.moe_ep_rank = 0
+        self.local_expert_mapping = None
+
+        if self.moe_ep_size > 1:
+            self.local_expert_mapping = torch.full((self.num_experts,), -1, dtype=torch.int32, device="cuda")
+            self.local_expert_mapping[
+                self.moe_ep_rank * self.num_local_experts : (self.moe_ep_rank + 1) * self.num_local_experts
+            ] = torch.arange(0, self.num_local_experts, dtype=torch.int32, device="cuda")
+
+    def dispatch(self, topk_ids) -> torch.Tensor:
+        if self.local_expert_mapping is not None:
+            return self.local_expert_mapping[topk_ids]
+        return topk_ids
+
+
+class Qwen3MoeSparseMoeBlock(nn.Module):
+    dispatcher = None
+    runner = None
+
+    def __init__(self, config):
+        super().__init__()
+        self.num_experts = config.num_experts
+        self.top_k = config.num_experts_per_tok
+        self.norm_topk_prob = config.norm_topk_prob
+
+        # gating
+        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
+
+        self.experts = nn.ModuleList(
+            [Qwen3MoeMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
+        )
+
+        if Qwen3MoeSparseMoeBlock.dispatcher is None:
+            Qwen3MoeSparseMoeBlock.dispatcher = StandardDispatcher(
+                num_experts=config.num_experts, num_local_experts=config.num_experts
+            )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        # router_logits: (batch * sequence_length, n_experts)
+        router_logits = self.gate(hidden_states)
+
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+
+        if self.norm_topk_prob:  # only diff with mixtral sparse moe block!
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        # we cast back to the input dtype
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        selected_experts = Qwen3MoeSparseMoeBlock.dispatcher.dispatch(selected_experts)
+
+        w13_weight = torch.stack(
+            [torch.cat([layer.gate_proj.weight, layer.up_proj.weight], dim=0) for layer in self.experts]
+        )
+        w2_weight = torch.stack([layer.down_proj.weight for layer in self.experts], dim=0)
+
+        final_hidden_states = fused_experts_impl(
+            hidden_states.to(torch.bfloat16),
+            w13_weight,
+            w2_weight,
+            routing_weights,
+            selected_experts,
+        )
+
+        return final_hidden_states, router_logits
+
+
+def apply_true_on_policy_patch_for_qwen3_moe():
+    from transformers.models.qwen3_moe import modeling_qwen3_moe
+
+    modeling_qwen3_moe.Qwen3MoeSparseMoeBlock = Qwen3MoeSparseMoeBlock
diff --git a/slime/backends/fsdp_utils/models/qwen3_moe_hf.py b/slime/backends/fsdp_utils/models/qwen3_moe_hf.py
new file mode 100644
index 0000000000..75c2286f19
--- /dev/null
+++ b/slime/backends/fsdp_utils/models/qwen3_moe_hf.py
@@ -0,0 +1,43 @@
+import torch
+import torch.nn.functional as F
+
+
+def apply_fsdp_moe_patch():
+
+    from transformers.models.qwen3_moe import modeling_qwen3_moe
+
+    def _forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        batch_size, sequence_length, hidden_dim = hidden_states.shape
+        hidden_states = hidden_states.view(-1, hidden_dim)
+        router_logits = self.gate(hidden_states)
+
+        routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)
+        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
+        if self.norm_topk_prob:
+            routing_weights /= routing_weights.sum(dim=-1, keepdim=True)
+        routing_weights = routing_weights.to(hidden_states.dtype)
+
+        final_hidden_states = torch.zeros(
+            (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
+        )
+
+        expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
+
+        # Loop over all experts
+        for expert_idx in range(self.num_experts):
+            expert_layer = self.experts[expert_idx]
+            idx, top_x = torch.where(expert_mask[expert_idx])
+
+            if top_x.numel() > 0:
+                current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
+                current_hidden_states = expert_layer(current_state) * routing_weights[top_x, idx, None]
+                final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
+            else:
+                # force experts to participate in computation graph
+                dummy_output = expert_layer(hidden_states[:1]) * 0.0
+                final_hidden_states[:1] = final_hidden_states[:1] + dummy_output
+
+        final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
+        return final_hidden_states, router_logits
+
+    modeling_qwen3_moe.Qwen3MoeSparseMoeBlock.forward = _forward
diff --git a/slime/backends/fsdp_utils/update_weight_utils.py b/slime/backends/fsdp_utils/update_weight_utils.py
new file mode 100644
index 0000000000..75e217beb9
--- /dev/null
+++ b/slime/backends/fsdp_utils/update_weight_utils.py
@@ -0,0 +1,286 @@
+import abc
+import logging
+import socket
+from argparse import Namespace
+from collections.abc import Sequence
+
+import ray
+import torch
+import torch.distributed as dist
+from ray.actor import ActorHandle
+from torch.distributed.tensor import DTensor, Replicate
+
+try:
+    from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions  # type: ignore[import]
+except ImportError:
+    from sglang.srt.patch_torch import monkey_patch_torch_reductions  # type: ignore[import]
+
+from sglang.srt.utils import MultiprocessingSerializer
+
+from slime.utils.distributed_utils import init_process_group
+
+
+try:
+    from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]
+except ImportError:
+    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket  # type: ignore[import]
+
+
+logger = logging.getLogger(__name__)
+
+
+class UpdateWeight(abc.ABC):
+    def __init__(self, args: Namespace, model: torch.nn.Module) -> None:
+        self.args = args
+        self.model = model
+        self.weight_version = 0
+
+    @abc.abstractmethod
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+        engine_gpu_counts: Sequence[int] | None = None,
+        engine_gpu_offsets: Sequence[int] | None = None,
+    ) -> None:
+        pass
+
+    def update_weights(self) -> None:
+        self.weight_version += 1
+        bucket = []
+        bucket_size = 0
+        for name, param in self.model.state_dict().items():
+            param_size = param.numel() * param.element_size()
+            if bucket and bucket_size + param_size >= self.args.update_weight_buffer_size:
+                self.wait_and_update_bucket_weights(bucket)
+                del bucket
+                bucket = []
+                bucket_size = 0
+
+            param = param.cuda()
+            if isinstance(param, DTensor):
+                # async version of param.full_tensor
+                param = param.redistribute(
+                    placements=[Replicate()] * param.device_mesh.ndim,
+                    async_op=True,
+                ).to_local()
+            bucket.append((name, param))
+            bucket_size += param_size
+
+        if bucket:
+            self.wait_and_update_bucket_weights(bucket)
+            del bucket
+            bucket = []
+            bucket_size = 0
+
+    def wait_and_update_bucket_weights(self, bucket):
+        bucket = [(name, param.wait()) if hasattr(param, "wait") else (name, param) for name, param in bucket]
+        self.update_bucket_weights(bucket, weight_version=self.weight_version)
+
+    @abc.abstractmethod
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        pass
+
+
+class UpdateWeightFromTensor(UpdateWeight):
+    """Push model weights to rollout engines using tensors.
+
+    Streams parameters in size-bounded buckets; optionally groups tensors by dtype
+    and flattens per dtype, gathers per-rank blobs to the source, and issues one
+    RPC per dtype per bucket (or one per bucket if not flattened).
+    """
+
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+        engine_gpu_counts: Sequence[int] | None = None,
+        engine_gpu_offsets: Sequence[int] | None = None,
+    ) -> None:
+        """Attach rollout engines and create per-engine IPC (Gloo) groups.
+
+        Sets the gather source rank, engine handle, and `tp_rank` within the
+        engine's local group.
+        """
+        self.rollout_engines = rollout_engines
+
+        if engine_gpu_counts is None:
+            engine_gpu_counts = [self.args.rollout_num_gpus_per_engine] * len(rollout_engines)
+        if engine_gpu_offsets is None:
+            # Fallback: assume engines are densely packed (no placeholder gaps).
+            engine_gpu_offsets = []
+            offset = 0
+            for c in engine_gpu_counts:
+                engine_gpu_offsets.append(offset)
+                offset += c
+
+        for i, engine in enumerate(self.rollout_engines):
+            start_rank = engine_gpu_offsets[i]
+            end_rank = start_rank + engine_gpu_counts[i]
+            group_ranks = list(range(start_rank, end_rank))
+            new_group = dist.new_group(
+                ranks=group_ranks,
+                backend="gloo",
+            )
+            if dist.get_rank() in group_ranks:
+                self._ipc_gather_src = start_rank
+                self._ipc_gather_group = new_group
+                self._ipc_engine = engine
+                # Calculate TP rank within this SGLang engine group
+                self.tp_rank = dist.get_rank() - start_rank
+
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        # Placeholder ranks (GPU slots reserved but no engine) have no gather group.
+        # gather_object is only collective among group members, so we skip entirely.
+        if self._ipc_gather_group is None:
+            return
+
+        monkey_patch_torch_reductions()
+        # Use flattened bucket approach similar to Megatron
+        logger.info("Using flattened tensor bucket")
+        # Group tensors by dtype (same as Megatron)
+        named_tensors_by_dtypes = {}
+        for name, tensor in named_tensors:
+            dtype = tensor.dtype
+            if dtype not in named_tensors_by_dtypes:
+                named_tensors_by_dtypes[dtype] = []
+            named_tensors_by_dtypes[dtype].append((name, tensor))
+
+        # Create flattened bucket for each dtype group
+        serialized_tensors = []
+        for _dtype, named_tensors in named_tensors_by_dtypes.items():
+            flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
+            metadata = flattened_tensor_bucket.get_metadata()
+            flattened_tensor_data = {
+                "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
+                "metadata": metadata,
+            }
+            serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
+
+        if self._ipc_gather_src == dist.get_rank():
+            # On rank 0, prepare a list to hold the gathered batches from all ranks.
+            gathered_serialized_batches = [None for _ in range(dist.get_world_size(self._ipc_gather_group))]
+        else:
+            gathered_serialized_batches = None
+
+        # Gather the serialized batches from all ranks to rank 0.
+        dist.gather_object(
+            obj=serialized_tensors,
+            object_gather_list=gathered_serialized_batches,
+            dst=self._ipc_gather_src,
+            group=self._ipc_gather_group,
+        )
+
+        if dist.get_rank() == self._ipc_gather_src:
+            # Handle flattened bucket format (same as Megatron approach)
+            # Each rank may have multiple dtype buckets
+            # TODO: here we assume all ranks have the same number of dtypes
+            num_dtypes = len(gathered_serialized_batches[0])
+            assert num_dtypes > 0
+            for i in range(num_dtypes):
+                kwargs = {
+                    "serialized_named_tensors": [tensors[i] for tensors in gathered_serialized_batches],
+                    "load_format": "flattened_bucket",
+                    "flush_cache": False,
+                    "weight_version": str(weight_version),
+                }
+                ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs)
+                ray.get(ref)
+
+        if dist.get_rank() == self._ipc_gather_src:
+            ref = self._ipc_engine.flush_cache.remote()
+            ray.get(ref)
+
+
+class UpdateWeightFromDistributed(UpdateWeight):
+    """Broadcast weights via a temporary NCCL group to rollout engines."""
+
+    def connect_rollout_engines(
+        self,
+        rollout_engines: Sequence[ActorHandle],
+        rollout_engine_lock: ActorHandle | None,
+        engine_gpu_counts: Sequence[int] | None = None,
+        engine_gpu_offsets: Sequence[int] | None = None,
+    ) -> None:
+        """On rank 0, initialize a temporary NCCL group for parameter broadcast."""
+        self.rollout_engines = rollout_engines
+        self.rollout_engine_lock = rollout_engine_lock
+
+        # For TP:
+        #   1. AllGather parameters to rank 0
+        #   2. Broadcast parameters from rank 0 to all sglang engines
+        self._is_src_rank = dist.get_rank() == 0
+
+        if engine_gpu_counts is None:
+            engine_gpu_counts = [self.args.rollout_num_gpus_per_engine] * len(rollout_engines)
+
+        if self._is_src_rank:
+            self._group_name = "slime"
+            master_address = ray._private.services.get_node_ip_address()
+            with socket.socket() as sock:
+                sock.bind(("", 0))
+                master_port = sock.getsockname()[1]
+            world_size = sum(engine_gpu_counts) + 1
+
+            # Compute cumulative rank offsets.
+            cumulative = [0]
+            for c in engine_gpu_counts:
+                cumulative.append(cumulative[-1] + c)
+
+            refs = [
+                engine.init_weights_update_group.remote(
+                    master_address,
+                    master_port,
+                    cumulative[i] + 1,
+                    world_size,
+                    self._group_name,
+                    backend="nccl",
+                )
+                for i, engine in enumerate(self.rollout_engines)
+            ]
+            self._model_update_groups = init_process_group(
+                backend="nccl",
+                init_method=f"tcp://{master_address}:{master_port}",
+                world_size=world_size,
+                rank=0,
+                group_name=self._group_name,
+            )
+            ray.get(refs)
+
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
+        """Send names/dtypes/shapes metadata to engines, then broadcast tensors.
+
+        Ensures tensors are contiguous; when `world_size == 1`, converts DTensors
+        to full tensors prior to `dist.broadcast`.
+        """
+        if not self._is_src_rank or not named_tensors:
+            return
+
+        refs = [
+            engine.update_weights_from_distributed.remote(
+                names=[name for name, _ in named_tensors],
+                dtypes=[param.dtype for _, param in named_tensors],
+                shapes=[param.shape for _, param in named_tensors],
+                group_name=self._group_name,
+                weight_version=str(weight_version),
+            )
+            for engine in self.rollout_engines
+        ]
+
+        handles = []
+        # Broadcast parameters one by one with memory management
+        for _name, param in named_tensors:
+            torch.cuda.empty_cache()
+            # Ensure tensor is contiguous and on the right device
+            param_data = param.data.contiguous()
+
+            # avoid `DTensor._op_dispatcher.dispatch` has `assert compute_mesh is not None` error
+            if dist.get_world_size() == 1 and isinstance(param_data, DTensor):
+                param_data = param_data.full_tensor()
+
+            # Synchronous broadcast to avoid memory buildup
+            handles.append(dist.broadcast(param_data, 0, group=self._model_update_groups, async_op=True))
+
+        for handle in handles:
+            handle.wait()
+        ray.get(refs)
diff --git a/slime/backends/megatron_utils/__init__.py b/slime/backends/megatron_utils/__init__.py
index b1936ae692..a4666fbeb9 100644
--- a/slime/backends/megatron_utils/__init__.py
+++ b/slime/backends/megatron_utils/__init__.py
@@ -40,5 +40,3 @@ def _patched_forward(self, *args, packed_seq_params=None, **kwargs):
     pass
 
 logging.getLogger("megatron").setLevel(logging.WARNING)
-
-from . import megatron_patch  # noqa: F401, E402
diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py
index 721a30c679..0091d0e9c3 100644
--- a/slime/backends/megatron_utils/actor.py
+++ b/slime/backends/megatron_utils/actor.py
@@ -1,10 +1,10 @@
 import logging
 import os
 import random
+import socket
 from argparse import Namespace
 from contextlib import nullcontext
 
-import numpy as np
 import ray
 import torch
 import torch.distributed as dist
@@ -13,19 +13,10 @@
 from torch_memory_saver import torch_memory_saver
 from transformers import AutoConfig, AutoTokenizer
 
-# torch_memory_saver is only used when the rollout backend is sglang (the
-# default). For the vLLM rollout backend we intentionally avoid importing it
-# (the preload .so is not loaded), so we resolve it lazily and tolerate the
-# import being unavailable.
-try:  # pragma: no cover - import guard
-    from torch_memory_saver import torch_memory_saver  # type: ignore
-except Exception:  # noqa: BLE001 - missing native lib should not crash import
-    torch_memory_saver = None  # type: ignore[assignment]
-
 from slime.ray.train_actor import TrainRayActor
 from slime.utils import train_dump_utils
 from slime.utils.data import process_rollout_data
-from slime.utils.distributed_utils import get_gloo_group
+from slime.utils.distributed_utils import get_gloo_group, init_process_group
 from slime.utils.logging_utils import init_tracking
 from slime.utils.memory_utils import clear_memory, print_memory
 from slime.utils.misc import Box
@@ -38,14 +29,13 @@
 from ...utils.tensor_backper import TensorBackuper
 from .checkpoint import load_checkpoint
 from .cp_utils import slice_log_prob_with_cp, slice_with_cp
-from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data
+from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data, sync_actor_critic_data
 from .initialize import init, is_megatron_main_rank
 from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values
 from .model import forward_only, initialize_model_and_optimizer, save, train
 from .update_weight.common import named_params_and_buffers
 from .update_weight.update_weight_from_distributed import UpdateWeightFromDistributed
 from .update_weight.update_weight_from_tensor import UpdateWeightFromTensor
-from .update_weight.update_weight_from_tensor_vllm import UpdateVLLMWeightFromTensor
 
 logging.getLogger("megatron").setLevel(logging.WARNING)
 
@@ -61,13 +51,13 @@ def init(
         with_ref: bool = False,
         with_opd_teacher: bool = False,
     ) -> int | None:
-        if args.debug_rollout_only:
-            self.args = args
-            return 0
-
         monkey_patch_torch_dist()
+
         super().init(args, role, with_ref, with_opd_teacher)
 
+        if args.debug_rollout_only:
+            return 0
+
         init(args)
 
         if is_megatron_main_rank():
@@ -88,17 +78,17 @@ def init(
         dist.barrier(group=get_gloo_group())
 
         if args.offload_train:
-            if self._uses_torch_memory_saver():
-                if (x := args.train_memory_margin_bytes) > 0:
-                    logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}")
-                    torch_memory_saver.memory_margin_bytes = x
-            else:
-                logger.info(
-                    "offload_train enabled with rollout_backend=vllm: using CPU offload "
-                    "instead of torch_memory_saver."
-                )
+            if (x := args.train_memory_margin_bytes) > 0:
+                logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}")
+                torch_memory_saver.memory_margin_bytes = x
+
+        if role == "critic":
+            self.args.load = self.args.critic_load
+            self.args.save = self.args.critic_save
+            self.args.lr = self.args.critic_lr
+            self.args.lr_warmup_iters = self.args.critic_lr_warmup_iters
 
-        self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id = initialize_model_and_optimizer(
+        (self.model, self.optimizer, self.opt_param_scheduler, loaded_rollout_id) = initialize_model_and_optimizer(
             args, role
         )
 
@@ -141,14 +131,7 @@ def init(
             hf_vocab = getattr(self.hf_config, "vocab_size", None)
             self.args.vocab_size = hf_vocab if hf_vocab is not None else self.tokenizer.vocab_size
 
-        use_vllm_colocate = self.args.colocate and getattr(self.args, "rollout_backend", "sglang") == "vllm"
-        use_tensor_update = self.args.colocate and getattr(self.args, "rollout_backend", "sglang") != "vllm"
-        if use_vllm_colocate:
-            update_weight_cls = UpdateVLLMWeightFromTensor
-        elif use_tensor_update:
-            update_weight_cls = UpdateWeightFromTensor
-        else:
-            update_weight_cls = UpdateWeightFromDistributed
+        update_weight_cls = UpdateWeightFromTensor if self.args.colocate else UpdateWeightFromDistributed
         self.weight_updater = update_weight_cls(
             self.args,
             self.model,
@@ -177,150 +160,15 @@ def init(
 
         return start_rollout_id
 
-    def _uses_torch_memory_saver(self) -> bool:
-        """Whether GPU memory should be released via torch_memory_saver.
-
-        torch_memory_saver relies on an LD_PRELOAD-injected allocator hook that
-        is only set up for the sglang rollout backend (see ``RayTrainGroup``).
-        For the vLLM rollout backend we fall back to an explicit CPU-offload
-        path that moves model parameters and optimizer state to CPU.
-        """
-        if not self.args.offload_train:
-            return False
-        if torch_memory_saver is None:
-            return False
-        return getattr(self.args, "rollout_backend", "sglang") != "vllm"
-
-    @torch.no_grad()
-    def _move_optimizer_state_to_device(self, device: str | torch.device) -> None:
-        """Move all tensor entries in optimizer state to ``device``.
-
-        Handles three layouts we see in practice:
-          * ``ChainedOptimizer`` used when there are multiple model chunks
-            (virtual pipeline) -- we must recurse into ``chained_optimizers``
-            because the outer object's ``.state`` is empty.
-          * Megatron ``DistributedOptimizer`` whose ``state`` is a ``ProxyDict``
-            that yields ``(key, tensor)`` directly (no ``.get``).
-          * Stock torch optimizer with ``dict[param, state_dict]`` layout.
-        """
-        if self.optimizer is None:
-            return
-
-        def _walk(opt) -> None:
-            # ChainedOptimizer (Megatron VP): dispatch to inner optimizers.
-            inner = getattr(opt, "chained_optimizers", None)
-            if inner:
-                for sub in inner:
-                    _walk(sub)
-                return
-
-            state = getattr(opt, "state", None)
-            if not state:
-                return
-            for key, value in list(state.items()):
-                if isinstance(value, torch.Tensor):
-                    state[key] = value.to(device, non_blocking=True)
-                elif isinstance(value, dict):
-                    for sub_key, sub_value in value.items():
-                        if isinstance(sub_value, torch.Tensor):
-                            value[sub_key] = sub_value.to(device, non_blocking=True)
-
-        _walk(self.optimizer)
-        torch.cuda.synchronize()
-
-    @torch.no_grad()
-    def _move_module_tensors_to_device(self, device: str | torch.device) -> None:
-        """Move params/grads/buffers to ``device`` in-place.
-
-        ``nn.Module.to()`` replaces ``nn.Parameter`` objects whenever the
-        device actually changes, which silently breaks any external references
-        the optimizer (or our ``weights_backuper``) still holds to the old
-        tensors -- gradients then accumulate on GPU tensors that the optimizer
-        no longer knows about, so ``optimizer.step()`` becomes a no-op on the
-        live parameters. Mutating ``p.data`` (and ``p.grad.data``) in place
-        keeps the ``Parameter`` identity intact, so every consumer continues
-        to see the moved storage.
-
-        Additional Megatron specifics handled here:
-          * ``p.main_grad`` (used with gradient-accumulation fusion / DDP) is
-            moved alongside ``p.grad``; the optimizer reads from
-            ``main_grad`` when present.
-          * Megatron ``DistributedDataParallel`` stores parameters/grads in a
-            few big contiguous ``ParamAndGradBuffer`` tensors; individual
-            ``p.data`` / ``p.main_grad`` are views into them. Moving only the
-            views leaves the big buckets allocated on GPU and leaves the
-            views pointing at stale storage, so we also migrate each
-            buffer's underlying ``param_data`` / ``grad_data`` in place.
-        """
-        if self.model is None:
-            return
-
-        def _move_tensor_attr(owner, attr: str) -> None:
-            t = getattr(owner, attr, None)
-            if isinstance(t, torch.Tensor):
-                t.data = t.data.to(device, non_blocking=True)
-
-        for module in self.model:
-            # 1) Move DDP bucket buffers first so subsequent param views land
-            #    on the new device's storage without orphaning the buckets.
-            for bucket_attr in ("buffers", "expert_parallel_buffers"):
-                bucket_list = getattr(module, bucket_attr, None)
-                # ``buffers`` on nn.Module is a *method*; only treat it as our
-                # target when it is the Megatron list of ParamAndGradBuffer.
-                if isinstance(bucket_list, (list, tuple)):
-                    for buf in bucket_list:
-                        for t_attr in ("param_data", "grad_data"):
-                            _move_tensor_attr(buf, t_attr)
-
-            # 2) Move parameter / grad views in place (Parameter identity
-            #    preserved so optimizer + weights_backuper stay consistent).
-            for p in module.parameters(recurse=True):
-                p.data = p.data.to(device, non_blocking=True)
-                if p.grad is not None:
-                    p.grad.data = p.grad.data.to(device, non_blocking=True)
-                _move_tensor_attr(p, "main_grad")
-
-            # 3) Registered buffers (e.g. RoPE caches, norm running stats).
-            #    Use the unbound ``nn.Module.buffers`` because Megatron's DDP
-            #    shadows the method with a ``list`` of ParamAndGradBuffer.
-            for b in torch.nn.Module.buffers(module, recurse=True):
-                b.data = b.data.to(device, non_blocking=True)
-
-    def _cpu_offload_sleep(self) -> None:
-        """CPU-offload alternative to ``torch_memory_saver.pause()`` for vLLM."""
-        self._move_module_tensors_to_device("cpu")
-        if getattr(self, "optimizer", None) is not None:
-            self._move_optimizer_state_to_device("cpu")
-        torch.cuda.synchronize()
-
-    def _cpu_offload_wake_up(self) -> None:
-        """CPU-offload counterpart to ``_cpu_offload_sleep``."""
-        device = torch.cuda.current_device()
-        self._move_module_tensors_to_device(device)
-        if getattr(self, "optimizer", None) is not None:
-            self._move_optimizer_state_to_device(device)
-        torch.cuda.synchronize()
-
     @timer
     def sleep(self) -> None:
         assert self.args.offload_train
 
         clear_memory(clear_host_memory=True)
         print_memory("before offload model")
-        if (
-            self.role == "actor"
-            and self.args.use_critic
-            and not self.args.colocate
-            and hasattr(self.weight_updater, "disconnect_rollout_engines")
-        ):
-            self.weight_updater.disconnect_rollout_engines()
         destroy_process_groups()
 
-        if self._uses_torch_memory_saver():
-            torch_memory_saver.pause()
-        else:
-            self._cpu_offload_sleep()
-            clear_memory()
+        torch_memory_saver.pause()
 
         print_memory("after offload model")
 
@@ -329,10 +177,7 @@ def wake_up(self) -> None:
         assert self.args.offload_train
         print_memory("before wake_up model")
 
-        if self._uses_torch_memory_saver():
-            torch_memory_saver.resume()
-        else:
-            self._cpu_offload_wake_up()
+        torch_memory_saver.resume()
 
         clear_memory()
         reload_process_groups()
@@ -359,14 +204,7 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch:
             # Move multimodal training tensors to GPU in advance
             rollout_data["multimodal_train_inputs"] = [
                 (
-                    {
-                        key: (
-                            torch.from_numpy(v.copy()).to(device=torch.cuda.current_device())
-                            if isinstance(v, np.ndarray)
-                            else v.to(device=torch.cuda.current_device())
-                        )
-                        for key, v in mm_dict.items()
-                    }
+                    {key: tensor.to(device=torch.cuda.current_device()) for key, tensor in mm_dict.items()}
                     if mm_dict is not None
                     else None
                 )
@@ -514,9 +352,9 @@ def compute_log_prob(
                 store_prefix=store_prefix,
             )
 
-    def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None):
+    def train(self, rollout_id: int, rollout_data_ref: Box) -> None:
         if self.args.debug_rollout_only:
-            return None
+            return
 
         if self.args.offload_train:
             self.wake_up()
@@ -525,22 +363,25 @@ def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None):
             rollout_data = self._get_rollout_data(rollout_data_ref)
 
         if self.role == "critic":
-            result = self.train_critic(rollout_id, rollout_data)
+            return self.train_critic(rollout_id, rollout_data)
         else:
-            self.train_actor(rollout_id, rollout_data, external_data=external_data)
-            result = None
-
-        if self.args.offload_train:
-            self.sleep()
+            return self.train_actor(rollout_id, rollout_data)
 
-        return result
-
-    def train_critic(self, rollout_id: int, rollout_data: RolloutBatch):
-        """Train critic and return CPU values (used as old-values for the next actor train)."""
+    def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
+        # Create data iterator for log_probs and train.
         data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
+        rollout_data.update(
+            forward_only(
+                get_values,
+                self.args,
+                self.model,
+                data_iterator,
+                num_microbatches,
+            )
+        )
 
-        # Compute current critic values (used as old_values for value loss and for actor advantages).
-        rollout_data.update(forward_only(get_values, self.args, self.model, data_iterator, num_microbatches))
+        if rollout_id >= self.args.num_critic_only_steps and not self.args.critic_train_only:
+            sync_actor_critic_data(self.args, rollout_data, self._actor_critic_groups)
 
         compute_advantages_and_returns(self.args, rollout_data)
 
@@ -554,13 +395,7 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch):
             num_microbatches,
         )
 
-        if mpu.is_pipeline_last_stage() and "values" in rollout_data:
-            from slime.backends.megatron_utils.data import tensors_to_cpu
-
-            return {"values": tensors_to_cpu(rollout_data["values"])}
-        return {}
-
-    def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data=None) -> None:
+    def train_actor(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
         # Create data iterator for log_probs and train.
         data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
 
@@ -595,21 +430,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
                     )
 
                 self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
-                can_reuse_log_probs_in_loss = (
-                    len(num_microbatches) == 1
-                    and self.args.loss_type == "policy_loss"
-                    and self.args.kl_coef == 0
-                    and not self.args.use_rollout_logprobs
-                    and not self.args.get_mismatch_metrics
-                    and not self.args.use_critic
-                    and not self.args.keep_old_actor
-                    and not self.args.use_opd
-                    and not self.args.use_routing_replay
-                    and self.args.advantage_estimator != "gspo"
-                )
-                if (
-                    not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics
-                ) and not can_reuse_log_probs_in_loss:
+                if not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics:
                     if self.args.use_routing_replay:
                         if self.args.use_rollout_routing_replay:
                             os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
@@ -626,12 +447,11 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
                         RoutingReplay.clear_all_forward()
 
                 if self.args.use_critic:
-                    if external_data is not None and mpu.is_pipeline_last_stage():
-                        values = external_data.get("values")
-                        if values is not None:
-                            from slime.backends.megatron_utils.data import tensors_to_gpu
-
-                            rollout_data["values"] = tensors_to_gpu(values)
+                    sync_actor_critic_data(
+                        self.args,
+                        rollout_data,
+                        self._actor_critic_groups,
+                    )
                 if self._active_model_tag != "actor":
                     self._switch_model("actor")
 
@@ -640,7 +460,7 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data
                 compute_advantages_and_returns(self.args, rollout_data)
 
             if self.rollout_data_postprocess is not None:
-                self.rollout_data_postprocess(self.args, rollout_id, rollout_data)
+                self.rollout_data_postprocess(self.args)
 
             log_rollout_data(
                 rollout_id,
@@ -691,7 +511,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
 
         # torch dist may trigger nccl communication during saving.
         if self.args.offload_train:
-            self.wake_up()
+            reload_process_groups()
 
         if self.args.async_save:
             from megatron.training.async_utils import maybe_finalize_async_save
@@ -709,7 +529,7 @@ def save_model(self, rollout_id: int, force_sync: bool = False) -> None:
             save_hf_model(self.args, rollout_id, self.model)
 
         if self.args.offload_train:
-            self.sleep()
+            destroy_process_groups()
 
     @timer
     def update_weights(self) -> None:
@@ -718,21 +538,17 @@ def update_weights(self) -> None:
 
         if self.args.use_fault_tolerance:
             if dist.get_rank() == 0:
-                ray.get(self.rollout_manager.recover_updatable_engines.remote())
+                ray.get(self.rollout_manager.recover_rollout_engines.remote())
             dist.barrier(group=get_gloo_group())
 
         rollout_engines, rollout_engine_lock, num_new_engines, engine_gpu_counts, engine_gpu_offsets = ray.get(
-            self.rollout_manager.get_updatable_engines_and_lock.remote()
+            self.rollout_manager.get_rollout_engines_and_lock.remote()
         )
 
-        reconnect_rollout_engines = self.args.offload_train and self.args.use_critic and not self.args.colocate
-
-        if reconnect_rollout_engines:
-            self.wake_up()
-        elif self.args.offload_train:
+        if self.args.offload_train:
             reload_process_groups()
 
-        if num_new_engines > 0 or reconnect_rollout_engines:
+        if num_new_engines > 0:
             self.weight_updater.connect_rollout_engines(
                 rollout_engines,
                 rollout_engine_lock,
@@ -741,11 +557,9 @@ def update_weights(self) -> None:
             )
             dist.barrier(group=get_gloo_group())
             if dist.get_rank() == 0:
-                ray.get(self.rollout_manager.clear_updatable_num_new_engines.remote())
+                ray.get(self.rollout_manager.clear_num_new_engines.remote())
 
-        weight_update_ctx = torch_memory_saver.disable() if self._uses_torch_memory_saver() else nullcontext()
-
-        with weight_update_ctx:
+        with torch_memory_saver.disable() if self.args.offload_train else nullcontext():
             print_memory("before update_weights")
             self.weight_updater.update_weights()
             print_memory("after update_weights")
@@ -769,9 +583,7 @@ def update_weights(self) -> None:
                 else:
                     self.weights_backuper.backup("old_actor")
 
-        if reconnect_rollout_engines:
-            self.sleep()
-        elif self.args.offload_train:
+        if self.args.offload_train:
             destroy_process_groups()
 
     def load_other_checkpoint(self, model_tag: str, path: str) -> None:
@@ -803,3 +615,26 @@ def load_other_checkpoint(self, model_tag: str, path: str) -> None:
 
         self.weights_backuper.backup(model_tag)
         self._active_model_tag = model_tag
+
+    def connect_actor_critic(
+        self,
+        actor_handle: ActorHandle | None = None,
+        master_address: str | None = None,
+        master_port: int | None = None,
+    ) -> None:
+        if self.role == "actor":
+            master_address = ray.util.get_node_ip_address()
+            with socket.socket() as sock:
+                sock.bind(("", 0))
+                master_port = sock.getsockname()[1]
+            actor_handle.connect_actor_critic.remote(master_address=master_address, master_port=master_port)
+
+        group_name = "actor_critic"
+        world_size = 2
+        self._actor_critic_groups = init_process_group(
+            backend="nccl",
+            init_method=f"tcp://{master_address}:{master_port}",
+            world_size=world_size,
+            rank=0 if self.role == "actor" else 1,
+            group_name=group_name,
+        )
diff --git a/slime/backends/megatron_utils/arguments.py b/slime/backends/megatron_utils/arguments.py
index 6a026271c8..d724b899ca 100644
--- a/slime/backends/megatron_utils/arguments.py
+++ b/slime/backends/megatron_utils/arguments.py
@@ -1,4 +1,3 @@
-import ast
 import logging
 
 from megatron.training.arguments import parse_args as _megatron_parse_args
@@ -11,67 +10,8 @@
 logger = logging.getLogger(__name__)
 
 
-_ALLGATHER_CP_DSA_ARCHITECTURES = {
-    "DeepseekV32ForCausalLM",
-    "GlmMoeDsaForCausalLM",
-}
-
-
-def _is_allgather_cp_dsa_model(hf_config):
-    if hf_config is None:
-        return False
-
-    architecture_names = getattr(hf_config, "architectures", None) or []
-    return any(name in _ALLGATHER_CP_DSA_ARCHITECTURES for name in architecture_names)
-
-
-def _validate_allgather_cp_supported(args, hf_config=None):
-    if not getattr(args, "allgather_cp", False) or getattr(args, "context_parallel_size", 1) <= 1:
-        return
-
-    if _is_allgather_cp_dsa_model(hf_config):
-        return
-
-    raise ValueError(
-        "--allgather-cp with --context-parallel-size > 1 is currently only supported for "
-        "DSA attention models (DeepSeek-V3.2 and GLM-5.1). Non-DSA models still use the "
-        "zigzag CP layout and would silently scramble token order under allgather CP. "
-        "Please remove --allgather-cp, set --context-parallel-size 1, or use a supported DSA model."
-    )
-
-
-def _has_dense_moe_layers(args):
-    moe_layer_freq = getattr(args, "moe_layer_freq", None)
-    if moe_layer_freq is None:
-        return True
-
-    if isinstance(moe_layer_freq, str):
-        try:
-            moe_layer_freq = ast.literal_eval(moe_layer_freq)
-        except (SyntaxError, ValueError):
-            return "0" in moe_layer_freq
-
-    try:
-        return any(int(layer_freq) == 0 for layer_freq in moe_layer_freq)
-    except TypeError:
-        return int(moe_layer_freq) == 0
-
-
-def _is_moe_config(hf_config):
-    return any(
-        hasattr(hf_config, attr)
-        for attr in (
-            "moe_intermediate_size",
-            "num_experts",
-            "n_routed_experts",
-            "num_local_experts",
-        )
-    )
-
-
 def validate_args(args):
     """Run megatron's own validate_args plus slime-specific megatron validations."""
-
     _megatron_validate_args(args)
 
     # always use varlen
@@ -100,52 +40,21 @@ def equal(x, y):
     if hasattr(hf_config, "text_config"):
         hf_config = hf_config.text_config
 
-    # Some models store rope_theta inside rope_parameters dict rather than
-    # as a top-level attribute.  Prefer the dict value when available so
-    # the validation doesn't compare against a stale class default.
-    rope_params = getattr(hf_config, "rope_parameters", None)
-    if isinstance(rope_params, dict) and "rope_theta" in rope_params:
-        _hf_rope_theta = rope_params["rope_theta"]
-    else:
-        _hf_rope_theta = getattr(hf_config, "rope_theta", None)
-
-    validate_dense_ffn = not _is_moe_config(hf_config) or _has_dense_moe_layers(args)
-
     for hf_config_name, megatron_config_name, compare_fn in [
         ("hidden_size", "hidden_size", equal),
         ("num_attention_heads", "num_attention_heads", equal),
         ("num_hidden_layers", "num_layers", equal),
         ("intermediate_size", "ffn_hidden_size", equal),
-        ("moe_intermediate_size", "moe_ffn_hidden_size", equal),
-        ("shared_expert_intermediate_size", "moe_shared_expert_intermediate_size", equal),
         ("tie_word_embeddings", "untie_embeddings_and_output_weights", lambda x, y: not x == y),
         ("rms_norm_eps", "norm_epsilon", equal),
-        ("rms_norm_eps", "layernorm_epsilon", equal),
+        ("rope_theta", "rotary_base", equal),
     ]:
         if hasattr(hf_config, hf_config_name):
-            if not hasattr(args, megatron_config_name):
-                logger.warning(
-                    f"Megatron args missing '{megatron_config_name}' (mapped from HF '{hf_config_name}') , "
-                    f"Skip validate"
+            if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)):
+                errors.append(
+                    f"{hf_config_name} in hf config {getattr(hf_config, hf_config_name)} is not equal to "
+                    f"{megatron_config_name} {getattr(args, megatron_config_name)}, please check the config."
                 )
-                continue
-        if hf_config_name == "intermediate_size" and not validate_dense_ffn:
-            continue
-
-        # if hasattr(hf_config, hf_config_name) and hasattr(args, megatron_config_name):
-        #     if not compare_fn(getattr(hf_config, hf_config_name), getattr(args, megatron_config_name)):
-        #         errors.append(
-        #             f"{hf_config_name} in hf config {getattr(hf_config, hf_config_name)} is not equal to "
-        #             f"{megatron_config_name} {getattr(args, megatron_config_name)}, please check the config."
-        #         )
-
-    # Validate rope_theta separately using the resolved value
-    if _hf_rope_theta is not None:
-        if not equal(_hf_rope_theta, getattr(args, "rotary_base", None)):
-            errors.append(
-                f"rope_theta in hf config {_hf_rope_theta} is not equal to "
-                f"rotary_base {getattr(args, 'rotary_base', None)}, please check the config."
-            )
 
     if len(errors) > 0:
         raise AssertionError("hf_validate_args failed: " + "; ".join(errors))
@@ -187,15 +96,14 @@ def megatron_parse_args(extra_args_provider, skip_hf_validate=False):
     """Parse megatron args, validate HF config, and set defaults."""
     args = _megatron_parse_args(extra_args_provider=extra_args_provider, ignore_unknown_args=True)
 
-    hf_config = None
     if args.hf_checkpoint and not skip_hf_validate:
         hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
         _hf_validate_args(args, hf_config)
 
-    if not skip_hf_validate:
-        _validate_allgather_cp_supported(args, hf_config)
-
     args.rank = 0
-    args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node
+    if args.critic_train_only:
+        args.world_size = args.critic_num_nodes * args.critic_num_gpus_per_node
+    else:
+        args.world_size = args.actor_num_nodes * args.actor_num_gpus_per_node
     args = _set_default_megatron_args(args)
     return args
diff --git a/slime/backends/megatron_utils/checkpoint.py b/slime/backends/megatron_utils/checkpoint.py
index d4eb236936..8c7cd5317c 100644
--- a/slime/backends/megatron_utils/checkpoint.py
+++ b/slime/backends/megatron_utils/checkpoint.py
@@ -135,9 +135,7 @@ 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 = megatron_bridge_utils.patch_auto_bridge_hf_config(
-            AutoBridge.from_hf_pretrained(load_path, trust_remote_code=True)
-        )
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, 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/cp_utils.py b/slime/backends/megatron_utils/cp_utils.py
index b0eab9b6ab..d30de602a8 100644
--- a/slime/backends/megatron_utils/cp_utils.py
+++ b/slime/backends/megatron_utils/cp_utils.py
@@ -223,10 +223,7 @@ def slice_log_prob_with_cp(
     qkv_format: str = "thd",
     max_token_len: int | None = None,
 ) -> list[float] | torch.Tensor:
-    assert len(log_prob) == response_length, (
-        f"log_prob length mismatch: len(log_prob)={len(log_prob)}, "
-        f"response_length={response_length}, total_length={total_length}"
-    )
+    assert len(log_prob) == response_length
 
     cp_size = mpu.get_context_parallel_world_size()
 
diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py
index 19db1f475a..36ef3dae99 100644
--- a/slime/backends/megatron_utils/data.py
+++ b/slime/backends/megatron_utils/data.py
@@ -72,7 +72,6 @@ def get_batch(
         assert max([t.size(0) for t in tokens]) <= max_seqlen
         tokens = [slice_with_cp(t, pad_token_id, qkv_format, max_seqlen) for t in tokens]
         tokens = torch.stack(tokens)
-        packed_seq_params = None
 
     elif qkv_format == "thd":
         if allgather_cp:
@@ -164,14 +163,18 @@ def get_batch(
     multimodal_train_inputs = batch.get("multimodal_train_inputs", None)
     if multimodal_train_inputs is not None:
         multimodal_data = {}  # key -> concatenated tensor
+        multimodal_num_items = {}  # key -> list of item counts per sequence
         for mm_input_dict in multimodal_train_inputs:
             if mm_input_dict is not None:
                 for key, mm_tensor in mm_input_dict.items():
                     if key not in multimodal_data:
                         multimodal_data[key] = mm_tensor
+                        multimodal_num_items[key] = [mm_tensor.size(0)]
                     else:
                         multimodal_data[key] = torch.cat([multimodal_data[key], mm_tensor], dim=0)
+                        multimodal_num_items[key].append(mm_tensor.size(0))
         batch["multimodal_train_inputs"] = multimodal_data
+        batch["multimodal_num_items"] = multimodal_num_items
 
     return batch
 
@@ -610,32 +613,53 @@ def log_perf_data(rollout_id: int, args: Namespace) -> None:
     )
 
 
-def tensors_to_cpu(tensor_list):
-    """Move a list of GPU tensors to CPU for Ray object store transfer.
-
-    Args:
-        tensor_list: List of GPU tensors, or None.
-
-    Returns:
-        List of CPU tensors (detached), or None if input is None.
+def sync_actor_critic_data(
+    args: Namespace,
+    rollout_data: RolloutBatch | None = None,
+    group: dist.ProcessGroup | None = None,
+) -> None:
     """
-    if tensor_list is None:
-        return None
-    return [t.detach().cpu() for t in tensor_list]
-
+    Broadcast `values` (from critic) and optionally `log_probs`/`ref_log_probs`
+    (from actor) across PP ranks to align data dependencies.
 
-def tensors_to_gpu(tensor_list, device=None):
-    """Move a list of CPU tensors back to GPU.
-
-    Args:
-        tensor_list: List of CPU tensors, or None.
-        device: Target CUDA device. If None, uses current device.
-
-    Returns:
-        List of GPU tensors, or None if input is None.
+    - Values are broadcast from src=1.
+    - Log-probs and ref-log-probs are broadcast from src=0 when KL is used.
+    Updates `rollout_data` in place with the synchronized tensors.
     """
-    if tensor_list is None:
-        return None
-    if device is None:
-        device = torch.cuda.current_device()
-    return [t.to(device=device, dtype=torch.float32) for t in tensor_list]
+    log_probs_key = "log_probs" if not args.use_rollout_logprobs else "rollout_log_probs"
+    values, log_probs, ref_log_probs = map(rollout_data.get, ("values", log_probs_key, "ref_log_probs"))
+
+    # return when not the pp last stage
+    if not values and not log_probs:
+        return
+
+    handles = []
+
+    if not values:
+        values = [torch.empty_like(log_prob) for log_prob in log_probs]
+    for value in values:
+        handles.append(dist.broadcast(value, src=1, group=group, async_op=True))
+
+    if args.kl_coef != 0 or args.use_kl_loss:
+        if not log_probs:
+            log_probs = [torch.empty_like(value) for value in values]
+        if not ref_log_probs:
+            ref_log_probs = [torch.empty_like(value) for value in values]
+        for ref_log_prob, log_prob in zip(ref_log_probs, log_probs, strict=False):
+            handles.append(dist.broadcast(log_prob, src=0, group=group, async_op=True))
+            handles.append(dist.broadcast(ref_log_prob, src=0, group=group, async_op=True))
+
+    for handle in handles:
+        handle.wait()
+
+    rollout_data.update(
+        {
+            k: v
+            for k, v in {
+                "values": values,
+                log_probs_key: log_probs,
+                "ref_log_probs": ref_log_probs,
+            }.items()
+            if v is not None
+        }
+    )
diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py
index 256338fca2..e84de361c2 100644
--- a/slime/backends/megatron_utils/loss.py
+++ b/slime/backends/megatron_utils/loss.py
@@ -145,7 +145,7 @@ def get_responses(
 def _allgather_cp_redistribute(
     res: dict[str, list[torch.Tensor]],
     *,
-    logits_local_len: int,
+    logits: torch.Tensor,
     args: Namespace,
     total_lengths: list[int],
     response_lengths: list[int],
@@ -162,7 +162,8 @@ def _allgather_cp_redistribute(
 
     Args:
         res: Dict mapping metric names to lists of per-sample tensors.
-        logits_local_len: Local sequence length on this rank.
+        logits: Model output used only to determine the local sequence length
+            (``logits.size(1)``).
         args: Configuration (needs ``qkv_format``).
         total_lengths: Total sequence lengths (prompt + response) per sample.
         response_lengths: Response segment lengths per sample.
@@ -170,19 +171,12 @@ def _allgather_cp_redistribute(
     """
     cp_group = mpu.get_context_parallel_group()
     cp_rank = mpu.get_context_parallel_rank()
+
+    logits_local_len = logits.size(1)  # logits shape: [1, T_local, ...]
     chunk_start = cp_rank * logits_local_len
     chunk_end = chunk_start + logits_local_len
 
     for key, values in res.items():
-        # Skip keys where all values are None (e.g. entropy when not computed)
-        if all(v is None for v in values):
-            continue
-
-        # Determine reference dtype/device from first non-None value
-        ref_value = next(v for v in values if v is not None)
-        ref_dtype = ref_value.dtype
-        ref_device = ref_value.device
-
         # Reconstruct full response tensors with each rank's contiguous contribution
         full_resps = []
         seq_start = 0
@@ -194,12 +188,12 @@ def _allgather_cp_redistribute(
             s = max(logit_global_start, chunk_start)
             e = min(logit_global_end, chunk_end)
 
-            if value is None or e <= s:
+            if e <= s:
                 # This rank has no response logprobs for this sample
                 full_resp = torch.zeros(
                     response_length,
-                    dtype=ref_dtype,
-                    device=ref_device,
+                    dtype=value.dtype,
+                    device=value.device,
                     requires_grad=True,
                 )
             else:
@@ -228,159 +222,6 @@ def _allgather_cp_redistribute(
         res[key] = new_values
 
 
-def _build_shifted_tokens(
-    T: int,
-    device: torch.device,
-    unconcat_tokens: list[torch.Tensor],
-    total_lengths: list[int],
-    response_lengths: list[int],
-    qkv_format: str,
-    max_seq_lens: list[int] | None,
-    allgather_cp: bool,
-) -> torch.Tensor:
-    """Build shifted target tokens for the full packed/padded logits."""
-    cp_size = mpu.get_context_parallel_world_size()
-
-    # --- zigzag CP: completely different layout ---
-    if cp_size > 1 and not allgather_cp:
-        full_tokens = torch.zeros(T, dtype=torch.long, device=device)
-        end = 0
-        for i, (tokens, total_length, response_length) in enumerate(
-            zip(unconcat_tokens, total_lengths, response_lengths, strict=False)
-        ):
-            max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None
-            chunk_size_cp, chunks_offset, logits_offset, tokens_offset = get_logits_and_tokens_offset_with_cp(
-                total_length, response_length, qkv_format, max_seq_len
-            )
-            for half, base in ((0, end), (1, end + chunk_size_cp)):
-                lo = logits_offset[half][0] - chunks_offset[half][0]
-                hi = logits_offset[half][1] - chunks_offset[half][0]
-                full_tokens[base + lo : base + hi] = tokens[tokens_offset[half][0] : tokens_offset[half][1]]
-            end += 2 * chunk_size_cp
-        return full_tokens
-
-    # --- cp1 and allgather-CP both build global shifted tokens the same way ---
-    T_global = sum(total_lengths) if allgather_cp else T
-    full_tokens = torch.zeros(T_global, dtype=torch.long, device=device)
-
-    if qkv_format == "thd" or allgather_cp:
-        offset = 0
-        for tokens, total_length in zip(unconcat_tokens, total_lengths, strict=False):
-            full_tokens[offset : offset + total_length - 1] = tokens[1:total_length]
-            offset += total_length
-    else:  # bshd, cp1
-        for i, (tokens, total_length) in enumerate(zip(unconcat_tokens, total_lengths, strict=False)):
-            seq_start = max_seq_lens[i] * i
-            full_tokens[seq_start : seq_start + total_length - 1] = tokens[1:total_length]
-
-    # allgather-CP: slice to local chunk
-    if allgather_cp:
-        cp_rank = mpu.get_context_parallel_rank()
-        chunk_start = cp_rank * T
-        chunk_end = chunk_start + T
-        if chunk_end <= T_global:
-            return full_tokens[chunk_start:chunk_end].contiguous()
-        local = torch.zeros(T, dtype=torch.long, device=device)
-        valid = T_global - chunk_start
-        if valid > 0:
-            local[:valid] = full_tokens[chunk_start:]
-        return local
-
-    return full_tokens
-
-
-def _extract_per_sample(
-    log_prob_full: torch.Tensor,
-    entropy_full: torch.Tensor | None,
-    total_lengths: list[int],
-    response_lengths: list[int],
-    qkv_format: str,
-    max_seq_lens: list[int] | None,
-    allgather_cp: bool,
-) -> tuple[list[torch.Tensor], list[torch.Tensor | None]]:
-    """Slice per-sample response log-probs/entropy from full-length 1-D tensors."""
-    cp_size = mpu.get_context_parallel_world_size()
-    log_probs_list: list[torch.Tensor] = []
-    entropy_list: list[torch.Tensor] = []
-
-    if cp_size > 1 and not allgather_cp:
-        # zigzag CP
-        pos = 0
-        for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)):
-            max_seq_len = max_seq_lens[i] if max_seq_lens is not None else None
-            chunk_size_cp, chunks_offset, logits_offset, _tokens_offset = get_logits_and_tokens_offset_with_cp(
-                total_length, response_length, qkv_format, max_seq_len
-            )
-            lo0 = logits_offset[0][0] - chunks_offset[0][0]
-            hi0 = logits_offset[0][1] - chunks_offset[0][0]
-            lo1 = logits_offset[1][0] - chunks_offset[1][0]
-            hi1 = logits_offset[1][1] - chunks_offset[1][0]
-
-            lp = torch.cat(
-                [
-                    log_prob_full[pos + lo0 : pos + hi0],
-                    log_prob_full[pos + chunk_size_cp + lo1 : pos + chunk_size_cp + hi1],
-                ],
-                dim=0,
-            )
-            log_probs_list.append(lp)
-            if entropy_full is not None:
-                ent = torch.cat(
-                    [
-                        entropy_full[pos + lo0 : pos + hi0],
-                        entropy_full[pos + chunk_size_cp + lo1 : pos + chunk_size_cp + hi1],
-                    ],
-                    dim=0,
-                )
-                entropy_list.append(ent)
-            pos += 2 * chunk_size_cp
-
-    elif allgather_cp:
-        cp_rank = mpu.get_context_parallel_rank()
-        local_len = log_prob_full.size(0)
-        chunk_start = cp_rank * local_len
-        chunk_end = chunk_start + local_len
-
-        seq_start = 0
-        for total_length, response_length in zip(total_lengths, response_lengths, strict=False):
-            prompt_length = total_length - response_length
-            logit_global_start = seq_start + prompt_length - 1
-            logit_global_end = seq_start + total_length - 1
-
-            s = max(logit_global_start, chunk_start)
-            e = min(logit_global_end, chunk_end)
-            if e <= s:
-                log_probs_list.append(torch.zeros((0,), dtype=log_prob_full.dtype, device=log_prob_full.device))
-                if entropy_full is not None:
-                    entropy_list.append(torch.zeros((0,), dtype=entropy_full.dtype, device=entropy_full.device))
-            else:
-                log_probs_list.append(log_prob_full[s - chunk_start : e - chunk_start])
-                if entropy_full is not None:
-                    entropy_list.append(entropy_full[s - chunk_start : e - chunk_start])
-            seq_start += total_length
-
-    else:
-        # cp1
-        if qkv_format == "thd":
-            offset = 0
-            for total_length, response_length in zip(total_lengths, response_lengths, strict=False):
-                end = offset + total_length
-                start = end - response_length
-                log_probs_list.append(log_prob_full[start - 1 : end - 1])
-                if entropy_full is not None:
-                    entropy_list.append(entropy_full[start - 1 : end - 1])
-                offset += total_length
-        else:  # bshd
-            for i, (total_length, response_length) in enumerate(zip(total_lengths, response_lengths, strict=False)):
-                end = max_seq_lens[i] * i + total_length
-                start = end - response_length
-                log_probs_list.append(log_prob_full[start - 1 : end - 1])
-                if entropy_full is not None:
-                    entropy_list.append(entropy_full[start - 1 : end - 1])
-
-    return log_probs_list, entropy_list
-
-
 def get_log_probs_and_entropy(
     logits: torch.Tensor,
     *,
@@ -394,63 +235,51 @@ def get_log_probs_and_entropy(
 ) -> dict[str, list[torch.Tensor]]:
     """Compute per-token log-probabilities (and optionally entropy) on responses.
 
-    Computes on the **full** logits ``[T, V]`` tensor at once (instead of
-    per-sample slicing) so backward traverses ``[T, V]`` only once, then
-    extracts per-sample response portions.
+    For each sample, extracts response-aligned logits and tokens, then computes
+    log-probabilities via softmax across the tensor-parallel group. Log-probs
+    are squeezed from `[R, 1]` to `[R]`. Entropy values are always appended
+    (even when `with_entropy=False`), but only included in the result dict
+    when requested.
+
+    Args:
+        logits: Policy logits with shape `[1, T, V]`.
+        args: Configuration (temperature applied in `get_responses`).
+        unconcat_tokens: List of token tensors per sample.
+        total_lengths: Total sequence lengths per sample.
+        response_lengths: Response segment lengths per sample.
+        with_entropy: If True, include "entropy" key in result.
+        non_loss_data: Unused; kept for API compatibility.
 
-    When ``entropy_coef == 0``, entropy is computed under ``torch.no_grad()``
-    to avoid retaining the computation graph and to skip cloning.
+    Returns:
+        Dict with key "log_probs" mapping to a list of `[R]` tensors per
+        sample. If `with_entropy` is True, also includes "entropy" key with
+        a list of `[R]` tensors.
     """
     assert non_loss_data
-    qkv_format = args.qkv_format
-
-    assert logits.dtype == torch.float32, f"{logits.dtype}"
-    assert len(logits.shape) == 3, f"{logits.shape}"
-
-    if qkv_format == "thd":
-        assert logits.size(0) == 1, f"{logits.shape}"
-        logits = logits.squeeze(0)
-    else:
-        assert max_seq_lens is not None
-        logits = logits.view(-1, logits.size(-1))
-
-    # Apply rollout temperature scaling to logits to match rollout-time log-probs.
-    rollout_temperature = getattr(args, "rollout_temperature", 1.0)
-    if rollout_temperature != 1.0:
-        logits = logits / rollout_temperature
-    logits = logits.contiguous()
-    T = logits.size(0)
-    device = logits.device
-    tp_group = mpu.get_tensor_model_parallel_group()
-    chunk_size = args.log_probs_chunk_size
-
-    # --- build full shifted-token target tensor ---
-    full_tokens = _build_shifted_tokens(
-        T, device, unconcat_tokens, total_lengths, response_lengths, qkv_format, max_seq_lens, args.allgather_cp
-    )
-
-    # --- compute on full [T,V] logits at once via calculate_log_probs_and_entropy ---
-    log_prob_full, entropy_full = calculate_log_probs_and_entropy(
+    log_probs_list = []
+    entropy_list = []
+    for logits_chunk, tokens_chunk in get_responses(
         logits,
-        full_tokens,
-        tp_group,
-        with_entropy=with_entropy,
-        chunk_size=chunk_size,
-    )
-    log_prob_full = log_prob_full.squeeze(-1)  # [T, 1] -> [T]
-
-    # --- extract per-sample response portions ---
-    log_probs_list, entropy_list = _extract_per_sample(
-        log_prob_full,
-        entropy_full,
-        total_lengths,
-        response_lengths,
-        qkv_format,
-        max_seq_lens,
-        args.allgather_cp,
-    )
+        args=args,
+        unconcat_tokens=unconcat_tokens,
+        total_lengths=total_lengths,
+        response_lengths=response_lengths,
+        max_seq_lens=max_seq_lens,
+    ):
+        log_prob, entropy = calculate_log_probs_and_entropy(
+            logits_chunk,
+            tokens_chunk,
+            mpu.get_tensor_model_parallel_group(),
+            with_entropy=with_entropy,
+            chunk_size=args.log_probs_chunk_size,
+        )
+
+        log_probs_list.append(log_prob.squeeze(-1))
+        entropy_list.append(entropy)
 
-    res = {"log_probs": log_probs_list}
+    res = {
+        "log_probs": log_probs_list,
+    }
     if with_entropy:
         res["entropy"] = entropy_list
 
@@ -458,14 +287,14 @@ def get_log_probs_and_entropy(
     if args.allgather_cp:
         _allgather_cp_redistribute(
             res,
-            logits_local_len=T,
+            logits=logits,
             args=args,
             total_lengths=total_lengths,
             response_lengths=response_lengths,
             max_seq_lens=max_seq_lens,
         )
 
-    return torch.empty((0,), device=device), res
+    return torch.empty((0,), device=logits.device), res
 
 
 def get_values(
@@ -517,7 +346,7 @@ def get_values(
     if args.allgather_cp:
         _allgather_cp_redistribute(
             res,
-            logits_local_len=logits.size(1),
+            logits=logits,
             args=args,
             total_lengths=total_lengths,
             response_lengths=response_lengths,
@@ -581,10 +410,6 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
     Early returns if both `log_probs` and `values` are None (intermediate
     pipeline stages).
 
-    If ``args.custom_advantage_function_path`` is set, it is called after KL computation
-    and must populate ``rollout_data["advantages"]`` and
-    ``rollout_data["returns"]``.
-
     Args:
         args: Configuration specifying estimator type, KL coefficient,
             normalization settings, and other hyperparameters.
@@ -593,10 +418,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
             "total_lengths"). Modified in-place to add "advantages" and
             "returns" keys, each mapping to lists of tensors per sample.
     """
-    rollout_log_probs: list[torch.Tensor] | None = rollout_data.get("rollout_log_probs")
-    log_probs: list[torch.Tensor] | None = (
-        rollout_log_probs if args.use_rollout_logprobs else rollout_data.get("log_probs")
-    )
+    log_probs: list[torch.Tensor] = rollout_data.get("rollout_log_probs" if args.use_rollout_logprobs else "log_probs")
     ref_log_probs: list[torch.Tensor] = rollout_data.get("ref_log_probs")
     rewards: list[float] = rollout_data.get("rewards")
     values: None | list[torch.Tensor] = rollout_data.get("values")
@@ -611,7 +433,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
 
     if args.kl_coef == 0 or not log_probs:
         # when kl_coef is 0, we won't compute ref_log_prob
-        xs = log_probs or rollout_log_probs or values
+        xs = log_probs if log_probs is not None else values
         kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
     else:
         kl = [
@@ -622,14 +444,8 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
             )
             for i in range(len(log_probs))
         ]
-    rollout_data["kl"] = kl
 
-    if args.custom_advantage_function_path is not None:
-        custom_adv_fn = load_function(args.custom_advantage_function_path)
-        custom_adv_fn(args, rollout_data)
-        advantages, returns = rollout_data["advantages"], rollout_data["returns"]
-
-    elif args.advantage_estimator in ["grpo", "gspo"]:
+    if args.advantage_estimator in ["grpo", "gspo"]:
         rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
         returns = get_grpo_returns(rewards, kl)
         # TODO: is the copy necessary?
@@ -825,7 +641,7 @@ def policy_loss_function(
         are enabled.
     """
     advantages = torch.cat(batch["advantages"], dim=0)
-    old_log_probs = batch["rollout_log_probs"] if args.use_rollout_logprobs else batch.get("log_probs")
+    old_log_probs = batch["rollout_log_probs"] if args.use_rollout_logprobs else batch["log_probs"]
 
     response_lengths = batch["response_lengths"]
     total_lengths = batch["total_lengths"]
@@ -842,11 +658,6 @@ def policy_loss_function(
     )
 
     log_probs = log_probs_and_entropy["log_probs"]
-    if not args.use_rollout_logprobs and not old_log_probs:
-        old_log_probs = [log_prob.detach() for log_prob in log_probs]
-    train_log_probs_for_tis = batch.get("log_probs")
-    if not train_log_probs_for_tis:
-        train_log_probs_for_tis = [log_prob.detach() for log_prob in log_probs]
 
     # Pre-gather log probs if needed by OPSM or GSPO to avoid duplicate gathering
     need_full_log_probs = args.use_opsm or args.advantage_estimator == "gspo"
@@ -915,7 +726,7 @@ def policy_loss_function(
         tis_kwargs = {
             "args": args,
             "pg_loss": pg_loss,
-            "train_log_probs": train_log_probs_for_tis,
+            "train_log_probs": batch["log_probs"],
             "rollout_log_probs": batch["rollout_log_probs"],
             "loss_masks": batch["loss_masks"],
             "total_lengths": total_lengths,
@@ -1183,18 +994,10 @@ def loss_function(
             raise ValueError(f"Unknown loss type: {args.loss_type}")
 
     if args.recompute_loss_function:
-        loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean, use_reentrant=False)
+        loss, log = checkpoint(func, args, batch, logits, sum_of_sample_mean)
     else:
         loss, log = func(args, batch, logits, sum_of_sample_mean)
 
-    # With allgather-CP, some CP ranks may have no loss-contributing tokens (e.g., all
-    # padding). Without this, gradient doesn't flow through their attention path, so
-    # the CP gather's backward (reduce-scatter) is not called, deadlocking other CP
-    # ranks that call it. Adding this zero loss forces autograd to traverse the full
-    # graph on every rank without changing gradient values.
-    if args.allgather_cp and mpu.get_context_parallel_world_size() > 1:
-        loss = loss + 0 * logits.sum()
-
     # Here we need to divide by cp_size because to cancel the multiply in Megatron.
     global_batch_size = batch.get("dynamic_global_batch_size", args.global_batch_size)
     if not args.calculate_per_token_loss:
diff --git a/slime/backends/megatron_utils/megatron_patch/__init__.py b/slime/backends/megatron_utils/megatron_patch/__init__.py
deleted file mode 100644
index 8693ff85a0..0000000000
--- a/slime/backends/megatron_utils/megatron_patch/__init__.py
+++ /dev/null
@@ -1 +0,0 @@
-from . import megatron_chunked_grad_coalesce_patch  # noqa: F401
diff --git a/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py b/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py
deleted file mode 100644
index f7ecceb2dd..0000000000
--- a/slime/backends/megatron_utils/megatron_patch/megatron_chunked_grad_coalesce_patch.py
+++ /dev/null
@@ -1,146 +0,0 @@
-# Patch _allreduce_non_tensor_model_parallel_grads (and its legacy alias
-# _allreduce_layernorm_grads) in megatron.core.distributed.finalize_model_grads
-# to coalesce/all_reduce TP-side grads in size-bounded chunks instead of one
-# large _flatten_dense_tensors(grads). Lowers the peak contiguous-memory
-# allocation during TP-side grad sync, avoiding OOM under allocator
-# fragmentation when the combined grad buffer would otherwise be very large.
-# SUM/AVG are element-wise, so chunking is mathematically equivalent.
-# Chunk size: SLIME_GRAD_COALESCE_CHUNK_BYTES, default 1 GiB.
-#
-# Cross-compatible across the Megatron-LM versions slime is run against:
-# the core_v0.13.0 line (DDP config exposes `use_custom_fsdp`, `_get_main_grad_attr`
-# takes `(param, use_custom_fsdp)`, target function takes `(model, config)`) and
-# the post-core_v0.15.0rc7 dev line (`use_megatron_fsdp`, single-arg
-# `_get_main_grad_attr`, `(model, config, tp_group)`). API differences are
-# resolved at runtime — no version-conditional imports.
-
-import inspect
-import logging
-import os
-import sys
-import warnings
-
-logger = logging.getLogger(__name__)
-
-try:
-    import torch
-    from megatron.core import parallel_state
-    from megatron.core.distributed.finalize_model_grads import (
-        _flatten_dense_tensors,
-        _get_main_grad_attr,
-        _reshard_if_dtensor,
-        _unflatten_dense_tensors,
-        _unshard_if_dtensor,
-        get_attr_wrapped_model,
-    )
-
-    # post-core_v0.15.0rc7 dev takes (param); core_v0.13.0 line takes
-    # (param, use_custom_fsdp=False).
-    _gma_takes_fsdp_arg = len(inspect.signature(_get_main_grad_attr).parameters) >= 2
-
-    def _grad_attr(param, fsdp_on):
-        if _gma_takes_fsdp_arg:
-            return _get_main_grad_attr(param, fsdp_on)
-        return _get_main_grad_attr(param)
-
-    def _fsdp_flag(ddp_config):
-        return bool(getattr(ddp_config, "use_megatron_fsdp", False) or getattr(ddp_config, "use_custom_fsdp", False))
-
-    _chunk_bytes = int(os.environ.get("SLIME_GRAD_COALESCE_CHUNK_BYTES") or (1 << 30))
-
-    def _split_into_chunks(params, grads, target_bytes):
-        """Greedy split keeping params/grads aligned. A single grad larger
-        than target_bytes is placed alone in its own chunk."""
-        chunks, cur_p, cur_g, cur_b = [], [], [], 0
-        for p, g in zip(params, grads, strict=False):
-            gb = g.numel() * g.element_size()
-            if cur_g and cur_b + gb > target_bytes:
-                chunks.append((cur_p, cur_g))
-                cur_p, cur_g, cur_b = [], [], 0
-            cur_p.append(p)
-            cur_g.append(g)
-            cur_b += gb
-        if cur_g:
-            chunks.append((cur_p, cur_g))
-        return chunks
-
-    def _allreduce_non_tensor_model_parallel_grads(model, config, tp_group=None):
-        # post-core_v0.15.0rc7 dev passes tp_group; core_v0.13.0 line omits it.
-        # Default-fill from parallel_state so the same body works for both call sites.
-        if tp_group is None:
-            tp_group = parallel_state.get_tensor_model_parallel_group()
-        if tp_group.size() <= 1:
-            return
-
-        params_sum, grads_sum = [], []
-        params_avg, grads_avg = [], []
-        ddp_config = None
-        for model_chunk in model:
-            ddp_config = model_chunk.ddp_config
-            fsdp_on = _fsdp_flag(ddp_config)
-            for name, param in get_attr_wrapped_model(model_chunk, "named_parameters")():
-                if not param.requires_grad:
-                    continue
-                if getattr(param, "average_gradients_across_tp_domain", False):
-                    target_params, target_grads = params_avg, grads_avg
-                elif (config.sequence_parallel and getattr(param, "sequence_parallel", False)) or (
-                    config.qk_layernorm and ("q_layernorm" in name or "k_layernorm" in name)
-                ):
-                    target_params, target_grads = params_sum, grads_sum
-                else:
-                    continue
-
-                grad_attr = _grad_attr(param, fsdp_on)
-                grad = getattr(param, grad_attr)
-                if grad is None:
-                    continue
-                target_params.append(param)
-                if fsdp_on and hasattr(grad, "_local_tensor"):
-                    target_grads.append(grad._local_tensor.data)
-                else:
-                    target_grads.append(_unshard_if_dtensor(grad).data)
-
-        for params, grads, op in (
-            (params_sum, grads_sum, torch.distributed.ReduceOp.SUM),
-            (params_avg, grads_avg, torch.distributed.ReduceOp.AVG),
-        ):
-            if not grads:
-                continue
-            fsdp_on = _fsdp_flag(ddp_config)
-            for p_chunk, g_chunk in _split_into_chunks(params, grads, _chunk_bytes):
-                coalesced = _flatten_dense_tensors(g_chunk)
-                torch.distributed.all_reduce(coalesced, op=op, group=tp_group)
-                for param, buf, synced in zip(
-                    p_chunk, g_chunk, _unflatten_dense_tensors(coalesced, g_chunk), strict=False
-                ):
-                    buf.copy_(synced)
-                    grad_attr = _grad_attr(param, fsdp_on)
-                    orig_grad = getattr(param, grad_attr)
-                    if fsdp_on and hasattr(orig_grad, "_local_tensor"):
-                        # buf already aliases orig_grad._local_tensor.data;
-                        # restore original DTensor wrapper (post-rc7 dev semantics).
-                        setattr(param, grad_attr, orig_grad)
-                    else:
-                        setattr(param, grad_attr, _reshard_if_dtensor(buf, orig_grad))
-                del coalesced
-
-    # The parent package re-exports a same-named function, shadowing the
-    # submodule attribute. Pull the real module out of sys.modules to setattr.
-    _fmg = sys.modules["megatron.core.distributed.finalize_model_grads"]
-    _fmg._allreduce_non_tensor_model_parallel_grads = _allreduce_non_tensor_model_parallel_grads
-    _fmg._allreduce_layernorm_grads = _allreduce_non_tensor_model_parallel_grads
-
-    logger.info(
-        "slime grad coalesce patch applied to "
-        "megatron.core.distributed.finalize_model_grads."
-        "_allreduce_non_tensor_model_parallel_grads (chunk=%d MiB)",
-        _chunk_bytes // (1 << 20),
-    )
-
-except ImportError as exc:
-    warnings.warn(
-        f"slime grad coalesce patch not applied — Megatron import failed ({exc!r}). "
-        "If this is a Megatron upgrade, the symbol layout may have changed; "
-        "without this patch, large-model TP grad sync may OOM.",
-        stacklevel=2,
-    )
diff --git a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py
index f97cfd6935..203f60eff0 100644
--- a/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py
+++ b/slime/backends/megatron_utils/megatron_to_hf/processors/quantizer_fp8.py
@@ -9,8 +9,7 @@
 
 def quantize_params_fp8(args, megatron_name, converted_named_params, quantization_config):
     assert quantization_config["quant_method"] == "fp8"
-    fmt = quantization_config.get("fmt", "e4m3")
-    assert fmt == "e4m3", f"Unsupported FP8 format: {fmt}"
+    assert quantization_config["fmt"] == "e4m3"
     assert quantization_config["activation_scheme"] == "dynamic"
     weight_block_size = quantization_config.get("weight_block_size", None)
 
@@ -76,10 +75,6 @@ def quantize_params_fp8(args, megatron_name, converted_named_params, quantizatio
         # indexer
         "self_attention.wq_b.weight",
         "self_attention.wk.weight",
-        # linear attention
-        "self_attention.linear_attn.in_proj_qkv.weight",
-        "self_attention.linear_attn.in_proj_z.weight",
-        "self_attention.linear_attn.out_proj.weight",
     ]:
         quantize_named_params = []
         for converted_name, param in converted_named_params:
diff --git a/slime/backends/megatron_utils/megatron_to_hf/qwen3_5.py b/slime/backends/megatron_utils/megatron_to_hf/qwen3_5.py
index 2aabd86eba..822760f8dd 100644
--- a/slime/backends/megatron_utils/megatron_to_hf/qwen3_5.py
+++ b/slime/backends/megatron_utils/megatron_to_hf/qwen3_5.py
@@ -12,7 +12,11 @@ def _convert_mtp_layer(args, name, param, layer_idx):
     if "final_layernorm.weight" in name:
         return [("mtp.norm.weight", param)]
     if "eh_proj.weight" in name:
-        return [("mtp.fc.weight", param)]
+        if param.dim() < 2:
+            raise ValueError(f"eh_proj weight expects 2D tensor, got {param.shape}")
+        first_half, second_half = param.chunk(2, dim=1)
+        new_param = torch.cat([second_half, first_half], dim=1)
+        return [("mtp.fc.weight", new_param)]
 
     if "transformer_layer" in name:
         proxy_name = name.replace(f"mtp.layers.{layer_idx}.transformer_layer", f"decoder.layers.{layer_idx}")
diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py
index f326b1d0d3..dd3724424d 100644
--- a/slime/backends/megatron_utils/model.py
+++ b/slime/backends/megatron_utils/model.py
@@ -18,10 +18,9 @@
 from megatron.core.optimizer.optimizer import MegatronOptimizer
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.pipeline_parallel import get_forward_backward_func
-from megatron.core.utils import get_model_config, unwrap_model
+from megatron.core.utils import get_model_config
 from megatron.training.global_vars import get_args
 from megatron.training.training import get_model
-from tqdm import tqdm
 
 from slime.utils import logging_utils
 from slime.utils.memory_utils import clear_memory
@@ -29,107 +28,11 @@
 from .checkpoint import load_checkpoint, save_checkpoint
 from .data import DataIterator, get_batch
 from .loss import loss_function
-from .model_provider import get_model_provider_func
+from .model_provider import get_model_provider_func, wrap_model_provider_with_freeze
 
 logger = logging.getLogger(__name__)
 
 
-def _disable_tqdm_for_non_main_rank() -> bool:
-    return not (
-        mpu.get_data_parallel_rank(with_context_parallel=True) == 0
-        and mpu.get_tensor_model_parallel_rank() == 0
-        and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1
-    )
-
-
-def _should_update_microbatch_pbar(model) -> bool:
-    if _disable_tqdm_for_non_main_rank():
-        return False
-
-    while hasattr(model, "module"):
-        model = model.module
-    vp_stage = getattr(model, "vp_stage", None)
-    if mpu.get_virtual_pipeline_model_parallel_world_size() is not None and vp_stage is not None:
-        return mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
-    return mpu.is_pipeline_last_stage(ignore_virtual=True)
-
-
-def _wrap_forward_step_with_microbatch_pbar(forward_step_func, pbar):
-    if pbar is None:
-        return forward_step_func
-
-    def wrapped_forward_step(*args, **kwargs):
-        result = forward_step_func(*args, **kwargs)
-        model = args[1] if len(args) > 1 else kwargs.get("model")
-        if model is not None and _should_update_microbatch_pbar(model):
-            pbar.update(1)
-        return result
-
-    return wrapped_forward_step
-
-
-def _iter_critic_output_layers(model: Sequence[DDP]):
-    for chunk_id, module in enumerate(unwrap_model(model)):
-        output_layer = getattr(module, "output_layer", None)
-        if output_layer is not None:
-            yield chunk_id, output_layer
-
-
-def _critic_output_layer_needs_reinit(args: Namespace, model: Sequence[DDP], role: str) -> bool:
-    if role != "critic" or args.load is None:
-        return False
-
-    from megatron.core.dist_checkpointing.serialization import load_tensors_metadata
-    from megatron.training.checkpointing import get_load_checkpoint_path_by_args
-
-    checkpoint_path = Path(get_load_checkpoint_path_by_args(args))
-    if not (checkpoint_path / ".metadata").is_file():
-        return False
-
-    checkpoint_metadata = load_tensors_metadata(str(checkpoint_path))
-    for _chunk_id, output_layer in _iter_critic_output_layers(model):
-        for name in ("weight", "bias"):
-            param = getattr(output_layer, name, None)
-            if param is None:
-                continue
-
-            param_name = f"output_layer.{name}"
-            ckpt_tensor_metadata = next(
-                (
-                    tensor_metadata
-                    for key, tensor_metadata in checkpoint_metadata.items()
-                    if key == param_name or key.endswith(f".{param_name}")
-                ),
-                None,
-            )
-            expected_shape = tuple(param.shape)
-            checkpoint_shape = tuple(ckpt_tensor_metadata.global_shape) if ckpt_tensor_metadata is not None else None
-            if checkpoint_shape == expected_shape:
-                continue
-
-            reason = (
-                "missing from checkpoint metadata"
-                if checkpoint_shape is None
-                else f"shape mismatch checkpoint={checkpoint_shape} runtime={expected_shape}"
-            )
-            logger.warning(
-                "Will reinitialize critic %s after checkpoint load because it is %s",
-                param_name,
-                reason,
-            )
-            return True
-
-    return False
-
-
-@torch.no_grad()
-def _reinitialize_critic_output_layer(model: Sequence[DDP]) -> None:
-    for _chunk_id, output_layer in _iter_critic_output_layers(model):
-        output_layer.weight.data.normal_(mean=0.0, std=0.02)
-        if output_layer.bias is not None:
-            output_layer.bias.data.zero_()
-
-
 def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler:
     """Create and configure the optimizer learning-rate/weight-decay scheduler.
 
@@ -202,7 +105,9 @@ def setup_model_and_optimizer(
     assert not args.moe_use_upcycling
     assert args.load is not None or args.pretrained_checkpoint is not None
 
-    model = get_model(get_model_provider_func(args, role), ModelType.encoder_or_decoder)
+    model = get_model(
+        wrap_model_provider_with_freeze(get_model_provider_func(args, role), args), ModelType.encoder_or_decoder
+    )
 
     # Optimizer
     kwargs = {}
@@ -319,17 +224,15 @@ def forward_step(
         packed_seq_params = batch["packed_seq_params"]
         total_lengths = batch["total_lengths"]
         response_lengths = batch["response_lengths"]
-        forward_kwargs = {
-            "input_ids": tokens,
-            "position_ids": None,
-            "attention_mask": None,
-            "labels": None,
-            "packed_seq_params": packed_seq_params,
-            "loss_mask": batch["full_loss_masks"],
-        }
-        if batch["multimodal_train_inputs"] is not None:
-            forward_kwargs.update(batch["multimodal_train_inputs"])
-        output_tensor = model(**forward_kwargs)
+        output_tensor = model(
+            input_ids=tokens,
+            position_ids=None,
+            attention_mask=None,
+            labels=None,
+            packed_seq_params=packed_seq_params,
+            loss_mask=batch["full_loss_masks"],
+            **(batch["multimodal_train_inputs"] if batch["multimodal_train_inputs"] is not None else {}),
+        )
 
         return output_tensor, partial(
             f,
@@ -356,18 +259,9 @@ def forward_step(
     config.timers = None
     forward_data_store = []
     num_steps_per_rollout = len(num_microbatches)
-    microbatch_pbar = tqdm(
-        total=sum(num_microbatches),
-        desc=f"{(store_prefix or getattr(model[0], 'role', 'actor')).rstrip('_')} forward",
-        unit="microbatch",
-        dynamic_ncols=True,
-        leave=False,
-        disable=_disable_tqdm_for_non_main_rank(),
-    )
-    forward_step_with_progress = _wrap_forward_step_with_microbatch_pbar(forward_step, microbatch_pbar)
     for step_id in range(num_steps_per_rollout):
         forward_data_store += forward_backward_func(
-            forward_step_func=forward_step_with_progress,
+            forward_step_func=forward_step,
             data_iterator=data_iterator,
             model=model,
             num_microbatches=num_microbatches[step_id],
@@ -375,7 +269,6 @@ def forward_step(
             micro_batch_size=args.micro_batch_size,
             forward_only=True,
         )
-    microbatch_pbar.close()
 
     # Move model back to the train mode.
     for model_module in model:
@@ -412,7 +305,6 @@ def train_one_step(
     optimizer: MegatronOptimizer,
     opt_param_scheduler: OptimizerParamScheduler,
     num_microbatches: int,
-    microbatch_pbar=None,
 ) -> tuple[dict[str, float], float]:
     """Execute a single pipeline-parallel training step.
 
@@ -492,10 +384,9 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
 
         if return_schedule_plan:
             assert not args.enable_mtp_training, "MTP training should not be enabled when using combined 1f1b"
-            position_ids = None
             output_tensor = model.build_schedule_plan(
                 input_ids=batch["tokens"],
-                position_ids=position_ids,
+                position_ids=None,
                 attention_mask=None,
                 labels=None,
                 packed_seq_params=batch["packed_seq_params"],
@@ -511,12 +402,12 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
                 "loss_mask": batch["full_loss_masks"],
             }
 
-            if batch["multimodal_train_inputs"] is not None:
-                forward_kwargs.update(batch["multimodal_train_inputs"])
-
             if args.enable_mtp_training:
                 forward_kwargs["mtp_kwargs"] = {"mtp_labels": batch["tokens"]}
 
+            if batch["multimodal_train_inputs"] is not None:
+                forward_kwargs.update(batch["multimodal_train_inputs"])
+
             output_tensor = model(**forward_kwargs)
 
         if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
@@ -527,7 +418,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
     # Forward pass.
     forward_backward_func = get_forward_backward_func()
     losses_reduced = forward_backward_func(
-        forward_step_func=_wrap_forward_step_with_microbatch_pbar(forward_step, microbatch_pbar),
+        forward_step_func=forward_step,
         data_iterator=data_iterator,
         model=model,
         num_microbatches=num_microbatches,
@@ -538,7 +429,6 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p
     )
 
     valid_step = True
-    grad_norm = float("nan")
     if not getattr(args, "check_for_nan_in_loss_and_grad", True):
         found_inf_flag = optimizer.prepare_grads()
         if found_inf_flag:
@@ -691,14 +581,6 @@ def train(
         pre_hook_enabled = False
 
     num_steps_per_rollout = len(num_microbatches)
-    microbatch_pbar = tqdm(
-        total=sum(num_microbatches),
-        desc=f"{getattr(model[0], 'role', 'actor')} train",
-        unit="microbatch",
-        dynamic_ncols=True,
-        leave=False,
-        disable=_disable_tqdm_for_non_main_rank(),
-    )
 
     # Run training iterations till done.
     for step_id in range(num_steps_per_rollout):
@@ -713,7 +595,6 @@ def train(
             optimizer,
             opt_param_scheduler,
             num_microbatches[step_id],
-            microbatch_pbar=microbatch_pbar,
         )
 
         if step_id == 0:
@@ -774,7 +655,7 @@ def train(
                     # TODO: figure out why KL is not exactly zero when using PPO loss with KL clipping, and whether this is expected behavior or a bug.
                     assert log_dict["train/ppo_kl"] < 1e-8, f"{log_dict=}"
                 if accumulated_step_id == 0 and "train/kl_loss" in log_dict:
-                    assert log_dict["train/kl_loss"] < 1e-8, f"{log_dict=}"
+                    assert log_dict["train/kl_loss"] == 0.0, f"{log_dict=}"
 
             logger.info(f"{role_tag}step {accumulated_step_id}: {log_dict}")
 
@@ -798,7 +679,6 @@ def train(
                     rel_tol=0.01,
                     abs_tol=0.01,
                 ), f"grad norm mismatch: {grad_norm} != {expected_grad_norm}"
-    microbatch_pbar.close()
     # Close out pre-hooks if using distributed optimizer and overlapped param gather.
     if pre_hook_enabled:
         disable_forward_pre_hook(model)
@@ -845,15 +725,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_auto_bridge_hf_config, patch_megatron_model
+        from slime.utils.megatron_bridge_utils import 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 = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True))
+        bridge = AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True)
 
         path.mkdir(parents=True, exist_ok=True)
 
@@ -886,7 +765,6 @@ def initialize_model_and_optimizer(
 
     if torch.version.hip:
         import megatron.core.dist_checkpointing.strategies.filesystem_async as filesystem_async_module
-
         from slime.utils.rocm_checkpoint_writer import ROCmFileSystemWriterAsync
 
         filesystem_async_module.FileSystemWriterAsync = ROCmFileSystemWriterAsync
@@ -894,7 +772,6 @@ def initialize_model_and_optimizer(
 
     model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role)
     model[0].role = role
-    reinit_critic_output_layer = _critic_output_layer_needs_reinit(args, model, role)
     clear_memory()
     iteration, _ = load_checkpoint(
         model,
@@ -903,10 +780,8 @@ def initialize_model_and_optimizer(
         checkpointing_context={},
         skip_load_to_model_and_opt=False,
     )
-    if reinit_critic_output_layer:
-        _reinitialize_critic_output_layer(model)
-        if (args.fp16 or args.bf16) and optimizer is not None:
-            optimizer.reload_model_params()
     clear_memory()
 
+    opt_param_scheduler.step(increment=iteration * args.global_batch_size)
+
     return model, optimizer, opt_param_scheduler, iteration
diff --git a/slime/backends/megatron_utils/model_provider.py b/slime/backends/megatron_utils/model_provider.py
index 68e62bb8d7..31db8b0da8 100644
--- a/slime/backends/megatron_utils/model_provider.py
+++ b/slime/backends/megatron_utils/model_provider.py
@@ -17,7 +17,6 @@
 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
 
 
@@ -55,7 +54,7 @@ def forward(
         return logits, None
 
 
-def _get_model_provider_func(
+def get_model_provider_func(
     args: argparse.Namespace,
     role: Literal["actor", "critic"] = "actor",
 ):
@@ -84,9 +83,7 @@ def wrapped_model_provider(
     if args.megatron_to_hf_mode == "bridge":
         from megatron.bridge import AutoBridge
 
-        import slime_plugins.megatron_bridge  # noqa: F401  # register custom bridges
-
-        bridge = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.hf_checkpoint, trust_remote_code=True))
+        bridge = 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
@@ -94,29 +91,11 @@ def wrapped_model_provider(
         provider.expert_model_parallel_size = args.expert_model_parallel_size
         provider.expert_tensor_parallel_size = args.expert_tensor_parallel_size
         provider.sequence_parallel = args.sequence_parallel
-        provider.context_parallel_size = args.context_parallel_size
-        provider.variable_seq_lengths = args.variable_seq_lengths
-        if hasattr(args, "moe_token_dispatcher_type"):
-            provider.moe_token_dispatcher_type = args.moe_token_dispatcher_type
         if getattr(args, "decoder_first_pipeline_num_layers", None) is not None:
             provider.num_layers_in_first_pipeline_stage = args.decoder_first_pipeline_num_layers
         if getattr(args, "decoder_last_pipeline_num_layers", None) is not None:
             provider.num_layers_in_last_pipeline_stage = args.decoder_last_pipeline_num_layers
         provider.finalize()
-
-        if role == "critic":
-            _original_provide = provider.provide
-
-            def _critic_provide(pre_process=True, post_process=True, vp_stage=None):
-                model = _original_provide(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
-                if post_process:
-                    model.output_layer = LinearForLastLayer(
-                        input_size=model.config.hidden_size, output_size=1, config=model.config
-                    )
-                return model
-
-            return _critic_provide
-
         return provider.provide
 
     def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage: int | None = None) -> GPTModel:
@@ -141,17 +120,7 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage
             transformer_layer_spec = import_module(args.spec)
             # Allow the spec to be a function so that user can use customized Megatron easier.
             if callable(transformer_layer_spec):
-                result = transformer_layer_spec(args, config, vp_stage)
-                # If the result is itself a model provider (callable with pre_process param),
-                # delegate model construction to it directly (e.g. glm-omni VL model).
-                if callable(result) and "pre_process" in inspect.signature(result).parameters:
-                    model = result(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
-                    if post_process and role == "critic":
-                        model.output_layer = LinearForLastLayer(
-                            input_size=config.hidden_size, output_size=1, config=config
-                        )
-                    return model
-                transformer_layer_spec = result
+                transformer_layer_spec = transformer_layer_spec(args, config, vp_stage)
         else:
             if args.num_experts:
                 # Define the decoder block spec
@@ -240,21 +209,13 @@ def model_provider(pre_process: bool = True, post_process: bool = True, vp_stage
 
 
 def wrap_model_provider_with_freeze(original_provider, args):
-    def wrapped_provider(
-        pre_process=True,
-        post_process=True,
-        **kwargs,
-    ):
+    def wrapped_provider(pre_process=True, post_process=True, vp_stage=None):
         sig = inspect.signature(original_provider)
-        provider_kwargs = {
-            "pre_process": pre_process,
-            "post_process": post_process,
-        }
-        for key in ["vp_stage", "config", "pg_collection"]:
-            if key in sig.parameters:
-                provider_kwargs[key] = kwargs.get(key, None)
+        if "vp_stage" in sig.parameters:
+            model = original_provider(pre_process=pre_process, post_process=post_process, vp_stage=vp_stage)
+        else:
+            model = original_provider(pre_process=pre_process, post_process=post_process)
 
-        model = original_provider(**provider_kwargs)
         freeze_model_params(model, args)
 
         return model
@@ -262,12 +223,8 @@ def wrapped_provider(
     return wrapped_provider
 
 
-def get_model_provider_func(args, role="actor"):
-    return wrap_model_provider_with_freeze(_get_model_provider_func(args, role), args)
-
-
 def freeze_model_params(model: GPTModel, args: argparse.Namespace):
-    if getattr(args, "only_train_params_name_list", None):
+    if args.only_train_params_name_list:
         for name, param in model.named_parameters():
             param.requires_grad = False
             for pattern in args.only_train_params_name_list:
@@ -275,7 +232,7 @@ def freeze_model_params(model: GPTModel, args: argparse.Namespace):
                     param.requires_grad = True
                     break
 
-    if getattr(args, "freeze_params_name_list", None):
+    if args.freeze_params_name_list:
         for name, param in model.named_parameters():
             for pattern in args.freeze_params_name_list:
                 if re.search(pattern, name):
diff --git a/slime/backends/megatron_utils/sglang.py b/slime/backends/megatron_utils/sglang.py
index 4c045c4817..97c82a31cd 100644
--- a/slime/backends/megatron_utils/sglang.py
+++ b/slime/backends/megatron_utils/sglang.py
@@ -1,9 +1,4 @@
 # the file to manage all sglang deps in the megatron actor
-# When sglang is installed we prefer its implementations; otherwise we
-# fall back to API-compatible reimplementations in
-# slime.backends.megatron_utils.weight_sync_utils.
-
-# ── FP8 quantisation helpers (sglang-only, no local fallback) ───────
 try:
     from sglang.srt.layers.quantization.fp8_utils import quant_weight_ue8m0, transform_scale_ue8m0
     from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
@@ -12,29 +7,19 @@
     transform_scale_ue8m0 = None
     should_deepgemm_weight_requant_ue8m0 = None
 
-# ── monkey_patch_torch_reductions ───────────────────────────────────
 try:
     from sglang.srt.utils.patch_torch import monkey_patch_torch_reductions
 except ImportError:
-    try:
-        from sglang.srt.patch_torch import monkey_patch_torch_reductions
-    except ImportError:
-        from .weight_sync_utils import monkey_patch_torch_reductions
+    from sglang.srt.patch_torch import monkey_patch_torch_reductions
+
+
+from sglang.srt.utils import MultiprocessingSerializer
 
-# ── MultiprocessingSerializer ───────────────────────────────────────
-try:
-    from sglang.srt.utils import MultiprocessingSerializer
-except ImportError:
-    from .weight_sync_utils import MultiprocessingSerializer
 
-# ── FlattenedTensorBucket ───────────────────────────────────────────
 try:
     from sglang.srt.weight_sync.tensor_bucket import FlattenedTensorBucket  # type: ignore[import]
 except ImportError:
-    try:
-        from sglang.srt.model_executor.model_runner import FlattenedTensorBucket  # type: ignore[import]
-    except ImportError:
-        from .weight_sync_utils import FlattenedTensorBucket
+    from sglang.srt.model_executor.model_runner import FlattenedTensorBucket  # type: ignore[import]
 
 __all__ = [
     "quant_weight_ue8m0",
diff --git a/slime/backends/megatron_utils/update_weight/common.py b/slime/backends/megatron_utils/update_weight/common.py
index d46fef298a..bac46519d8 100644
--- a/slime/backends/megatron_utils/update_weight/common.py
+++ b/slime/backends/megatron_utils/update_weight/common.py
@@ -12,38 +12,6 @@
 from slime.utils.types import ParamInfo
 
 
-def _cat_partitions(
-    partitions: list[torch.Tensor], partition_dim: int, stride: int
-) -> torch.Tensor:
-    """Concatenate TP-gathered partitions, reversing Megatron-Core's strided layout.
-
-    When ``stride == 1`` each rank holds a single contiguous shard and a plain
-    ``torch.cat`` along *partition_dim* is sufficient.
-
-    When ``stride > 1``, Megatron-Core's ``_initialize_affine_weight_cpu``
-    splits the master weight into ``world_size * stride`` equal sub-chunks and
-    assigns rank *r* the sub-chunks at indices ``r, r + world_size,
-    r + 2 * world_size, …`` (i.e. ``weight_list[rank::world_size]``).  Each
-    rank therefore holds ``stride`` sub-chunks concatenated together.
-
-    To reconstruct the original full tensor we split each rank's shard back
-    into ``stride`` sub-chunks and interleave them:
-    ``for i in range(stride): for r in range(world_size): append chunk[r][i]``
-    """
-    if stride == 1:
-        return torch.cat(partitions, dim=partition_dim)
-
-    world_size = len(partitions)
-    # Each rank holds exactly `stride` sub-chunks concatenated along partition_dim.
-    rank_chunks = [p.chunk(stride, dim=partition_dim) for p in partitions]
-    # Reconstruct original ordering: chunk index (i * world_size + r)
-    ordered: list[torch.Tensor] = []
-    for i in range(stride):
-        for r in range(world_size):
-            ordered.append(rank_chunks[r][i])
-    return torch.cat(ordered, dim=partition_dim)
-
-
 def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
     """
     All-gather TP-sharded param to full tensor. expert_bias→param, non-TP/duplicated→param.data.
@@ -66,23 +34,17 @@ def all_gather_param(name: str, param: torch.nn.Parameter) -> torch.Tensor:
     param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
     dist.all_gather(param_partitions, param.data, group=tp_group)
     partition_dim = param.partition_dim
-    stride = getattr(param, "partition_stride", 1)
-    assert param.partition_stride == 1 or (
-        param.partition_stride == 2 and "linear_fc1" in name
-    ), "partition_stride != 1 is not supported"
+    assert param.partition_stride == 1, "partition_stride != 1 is not supported"
     # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
     # TODO: check only GLU is used.
-    if "linear_fc1.weight" in name or "linear_fc1.bias" in name:
+    if "linear_fc1.weight" in name:
         param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
         param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-        # GLU rechunking already reverses the stride=2 interleaving, so use
-        # plain concatenation from here on.
-        stride = 1
     # this is bug in megatron's grouped moe.
     if "linear_fc2.weight" in name:
         if partition_dim == 0:
             partition_dim = 1
-    param = _cat_partitions(param_partitions, partition_dim, stride)
+    param = torch.cat(param_partitions, dim=partition_dim)
     return param
 
 
@@ -101,10 +63,10 @@ def all_gather_params_async(
     for info, param in param_infos_and_params:
         # Prepare async all_gather
         if "expert_bias" in info.name:
-            gather_tasks.append((info, param, None, None, None, 1))
+            gather_tasks.append((info, param, None, None, None))
             handles.append(None)
         elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
-            gather_tasks.append((info, param.data, None, None, None, 1))
+            gather_tasks.append((info, param.data, None, None, None))
             handles.append(None)
         else:
             # Start async all_gather
@@ -117,8 +79,7 @@ def all_gather_params_async(
 
             param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
             handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True)
-            partition_stride = getattr(param, "partition_stride", 1)
-            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim, partition_stride))
+            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim))
             handles.append(handle)
 
     # Phase 2: Wait for ALL async operations to complete at once
@@ -129,24 +90,23 @@ def all_gather_params_async(
 
     # Phase 3: Process all results after all communications are done
     gathered_params = []
-    for info, direct_param, handle, param_partitions, partition_dim, partition_stride in gather_tasks:
+    for info, direct_param, handle, param_partitions, partition_dim in gather_tasks:
         if handle is None:
             # No all_gather needed
             param = direct_param
         else:
             # Process the gathered partitions (same logic as original all_gather_param)
+            assert partition_dim is not None, "partition_stride != 1 is not supported"
             # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
             # TODO: check only GLU is used.
-            if "linear_fc1.weight" in info.name or "linear_fc1.bias" in info.name:
+            if "linear_fc1.weight" in info.name:
                 param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
                 param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-                # GLU rechunking already reverses the stride=2 interleaving.
-                partition_stride = 1
             # this is bug in megatron's grouped moe.
             if "linear_fc2.weight" in info.name:
                 if partition_dim == 0:
                     partition_dim = 1
-            param = _cat_partitions(param_partitions, partition_dim, partition_stride)
+            param = torch.cat(param_partitions, dim=partition_dim)
 
         gathered_params.append(param)
 
@@ -171,10 +131,7 @@ def named_params_and_buffers(
 
 
 def _maybe_get_cpu_backup(x: torch.Tensor):
-    try:
-        from torch_memory_saver import torch_memory_saver
-    except Exception:  # noqa: BLE001 - native lib may be unavailable (e.g. vLLM backend)
-        return x
+    from torch_memory_saver import torch_memory_saver
 
     if (cpu_tensor := torch_memory_saver.get_cpu_backup(x)) is not None:
         return cpu_tensor
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 9ea728697b..01266aacef 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,9 +44,7 @@ def __init__(self, *args, **kwargs):
 
         import slime_plugins.megatron_bridge  # noqa: F401
 
-        self._bridge = megatron_bridge_utils.patch_auto_bridge_hf_config(
-            AutoBridge.from_hf_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
-        )
+        self._bridge = 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):
@@ -82,8 +80,6 @@ def _streaming_quantized():
 
 def _process_conversion_tasks(vanilla_conversion_tasks, new_weight_dict):
     def _handle_one(task):
-        if task is None:
-            return None
         if task.param_weight is None:
             return task
 
diff --git a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
index cd1976f90c..a353e353fe 100644
--- a/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
+++ b/slime/backends/megatron_utils/update_weight/hf_weight_iterator_direct.py
@@ -170,6 +170,7 @@ def _get_megatron_local_param_infos(args: Namespace, model: Sequence[torch.nn.Mo
                 continue
             for name, info in infos.items():
                 if name in param_infos:
+                    assert args.mtp_num_layers is not None
                     old_info = param_infos[name]
                     if old_info.src_rank > src_rank:
                         param_infos[name] = info
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
index c504212f64..7b5b7817f1 100644
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_distributed.py
@@ -1,9 +1,5 @@
-import logging
-import torch.multiprocessing as mp
-import os
 import socket
 import time
-import traceback
 from argparse import Namespace
 from collections.abc import Callable, Mapping, Sequence
 
@@ -20,155 +16,6 @@
 from ..megatron_to_hf import convert_to_hf
 from .common import all_gather_param, named_params_and_buffers
 
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# NcclBridge: isolate vLLM's PyNcclCommunicator in a subprocess so that it
-# never coexists with torch.distributed NCCL groups in the Megatron trainer.
-#
-# vLLM's weight transfer uses raw NCCL (PyNcclCommunicator) which conflicts
-# with torch.distributed's NCCL backend when both exist in the same process
-# (see https://github.com/vllm-project/vllm/issues/5477).  sglang avoids
-# this because it uses torch.distributed process groups for weight sync.
-# ---------------------------------------------------------------------------
-
-
-def _nccl_bridge_worker(conn, master_address, master_port, world_size, device, cvd, env_snapshot):
-    """Subprocess entry-point: creates PyNcclCommunicator and serves requests.
-
-    GPU tensors are shared from the parent via CUDA IPC (torch.multiprocessing
-    handles this transparently).  No GPU→CPU→GPU copies are needed.
-
-    Protocol over *conn* (multiprocessing.Connection):
-        parent → child:
-            {"op": "broadcast", "tensors": [gpu_tensor, ...]}
-            {"op": "send_packed", "named_tensors": [(name, gpu_tensor), ...]}
-            None  → shutdown
-        child → parent:
-            "ready"  (after init)
-            "ok"     (after each op)
-            "error: <msg>"
-    """
-    try:
-        os.environ.update(env_snapshot)
-        if cvd:
-            os.environ["CUDA_VISIBLE_DEVICES"] = cvd
-
-        import torch
-        import torch.multiprocessing  # ensure CUDA IPC reducers are registered
-        torch.cuda.set_device(device)
-
-        from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
-        from vllm.distributed.utils import StatelessProcessGroup
-
-        pg = StatelessProcessGroup.create(
-            host=master_address, port=master_port, rank=0, world_size=world_size,
-        )
-        comm = PyNcclCommunicator(pg, device=device)
-
-        conn.send("ready")
-
-        while True:
-            cmd = conn.recv()
-            if cmd is None:
-                break
-
-            op = cmd["op"]
-            if op == "broadcast":
-                for t in cmd["tensors"]:
-                    comm.broadcast(t, src=0, stream=torch.cuda.current_stream())
-                torch.cuda.synchronize()
-                conn.send("ok")
-
-            elif op == "send_packed":
-                from vllm.distributed.weight_transfer.nccl_engine import (
-                    NCCLTrainerSendWeightsArgs,
-                    NCCLWeightTransferEngine,
-                )
-
-                trainer_args = NCCLTrainerSendWeightsArgs(
-                    group=comm,
-                    packed=True,
-                )
-                NCCLWeightTransferEngine.trainer_send_weights(
-                    iterator=iter(cmd["named_tensors"]),
-                    trainer_args=trainer_args,
-                )
-                torch.cuda.synchronize()
-                conn.send("ok")
-
-    except Exception as e:
-        try:
-            conn.send(f"error: {e}")
-        except Exception:
-            pass
-        traceback.print_exc()
-
-
-class _NcclBridge:
-    """Runs vLLM's PyNcclCommunicator in a separate subprocess.
-
-    This prevents NCCL communicator conflicts with torch.distributed groups
-    that already exist in the Megatron trainer process.  GPU tensors are shared
-    with the subprocess via CUDA IPC (handled transparently by
-    torch.multiprocessing), avoiding any GPU→CPU→GPU copies.
-    """
-
-    def __init__(self, master_address: str, master_port: int, world_size: int, device: int):
-        ctx = mp.get_context("spawn")
-        self._parent_conn, child_conn = ctx.Pipe()
-
-        # Pass the full environment so the subprocess inherits all NCCL, CUDA,
-        # and networking settings from the trainer (set by Ray runtime_env).
-        env_snapshot = dict(os.environ)
-        cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
-
-        self._process = ctx.Process(
-            target=_nccl_bridge_worker,
-            args=(child_conn, master_address, master_port, world_size, device, cvd, env_snapshot),
-            daemon=True,
-        )
-        self._process.start()
-
-        msg = self._parent_conn.recv()
-        if isinstance(msg, str) and msg.startswith("error:"):
-            raise RuntimeError(f"NcclBridge init failed: {msg}")
-        if msg != "ready":
-            raise RuntimeError(f"NcclBridge init unexpected response: {msg}")
-        logger.info("NcclBridge ready (pid=%d, device=%d)", self._process.pid, device)
-
-    def broadcast_tensors(self, tensors: list[torch.Tensor]) -> None:
-        """Broadcast a list of tensors (one-by-one) via the bridge subprocess."""
-        gpu_tensors = [t.contiguous() for t in tensors]
-        self._parent_conn.send({"op": "broadcast", "tensors": gpu_tensors})
-        self._wait_ok("broadcast_tensors")
-
-    def send_weights_packed(self, named_tensors: list[tuple[str, torch.Tensor]]) -> None:
-        """Send weights using vLLM's packed broadcast protocol."""
-        gpu_pairs = []
-        for name, t in named_tensors:
-            data = t.data if hasattr(t, "data") else t
-            gpu_pairs.append((name, data.contiguous()))
-        self._parent_conn.send({"op": "send_packed", "named_tensors": gpu_pairs})
-        self._wait_ok("send_weights_packed")
-
-    def _wait_ok(self, label: str, timeout: float = 600.0) -> None:
-        if not self._parent_conn.poll(timeout):
-            raise TimeoutError(f"NcclBridge {label} timed out after {timeout}s")
-        msg = self._parent_conn.recv()
-        if msg != "ok":
-            raise RuntimeError(f"NcclBridge {label} failed: {msg}")
-
-    def shutdown(self) -> None:
-        try:
-            self._parent_conn.send(None)
-            self._process.join(timeout=30)
-        except Exception:
-            pass
-        if self._process.is_alive():
-            self._process.terminate()
-
 
 class UpdateWeightFromDistributed:
     """
@@ -231,14 +78,6 @@ def connect_rollout_engines(
                 engine_gpu_counts=engine_gpu_counts,
             )
 
-    def disconnect_rollout_engines(self) -> None:
-        if not getattr(self, "_is_pp_src_rank", False) or self._model_update_groups is None:
-            return
-        disconnect_rollout_engines_from_distributed(
-            self.args, self._group_name, self._model_update_groups, self.rollout_engines
-        )
-        self._model_update_groups = None
-
     @torch.no_grad()
     def update_weights(self) -> None:
         """
@@ -259,52 +98,34 @@ def update_weights(self) -> None:
                 )
         dist.barrier(group=get_gloo_group())
 
-        use_vllm_packed = self._use_vllm_packed()
-        if use_vllm_packed and self._is_pp_src_rank:
-            logger.info("Using vLLM packed weight sync (one-shot metadata + trainer_send_weights)")
-        if use_vllm_packed:
-            # vLLM packed path: gather all non-expert params, one-shot update (aligned with vLLM New Weight Syncing)
-            converted_named_tensors = []
-            for name, param in named_params_and_buffers(self.args, self.model):
-                if ".experts." in name:
-                    continue
-                param = all_gather_param(name, param)
-                if self._is_pp_src_rank:
-                    converted_named_tensors += convert_to_hf(
-                        self.args, self.model_name, name, param, self.quantization_config
-                    )
-            if converted_named_tensors and self._is_pp_src_rank:
-                self._update_weights_vllm_packed(converted_named_tensors)
-        else:
-            buffer_size = 0
-            converted_named_tensors = []
-            pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None
-
-            for name, param in named_params_and_buffers(self.args, self.model):
-                if ".experts." in name:
-                    continue
-                buffer_size = self._update_weight_from_distributed(
-                    name, param, converted_named_tensors, buffer_size, pbar=pbar
-                )
+        buffer_size = 0
+        converted_named_tensors = []
+        # non expert params
+        pbar = tqdm(desc=f"[{self._group_name}] Update weights", total=0) if self._is_pp_src_rank else None
+
+        for name, param in named_params_and_buffers(self.args, self.model):
+            if ".experts." in name:
+                continue
+            buffer_size = self._update_weight_from_distributed(
+                name, param, converted_named_tensors, buffer_size, pbar=pbar
+            )
 
-            if converted_named_tensors:
-                self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
+        if converted_named_tensors:
+            self._update_bucket_weights_from_distributed(converted_named_tensors, pbar=pbar)
 
         dist.barrier(group=get_gloo_group())
 
-        if not use_vllm_packed:
-            buffer_size = 0
-            named_tensors = []
-            pbar = tqdm(desc=f"[{self._group_name}] Update weights (experts)", total=0) if self._is_pp_src_rank else None
-            for name, param in named_params_and_buffers(self.args, self.model):
-                if ".experts." not in name:
-                    continue
-                buffer_size = self._update_expert_weight_from_distributed(
-                    name, param, named_tensors, buffer_size, pbar=pbar
-                )
+        buffer_size = 0
+        named_tensors = []
+        for name, param in named_params_and_buffers(self.args, self.model):
+            if ".experts." not in name:
+                continue
+            buffer_size = self._update_expert_weight_from_distributed(
+                name, param, named_tensors, buffer_size, pbar=pbar
+            )
 
-            if named_tensors:
-                self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
+        if named_tensors:
+            self._update_expert_bucket_weights_from_distributed(named_tensors, pbar=pbar)
 
         dist.barrier(group=get_gloo_group())
         if dist.get_rank() == 0:
@@ -342,41 +163,6 @@ def _update_weight_from_distributed(
         buffer_size += param_size
         return buffer_size
 
-    def _use_vllm_packed(self) -> bool:
-        """Use vLLM packed weight transfer (one-shot metadata + trainer_send_weights)."""
-        if not _is_vllm_backend(self.args):
-            return False
-        if not getattr(self.args, "vllm_weight_sync_packed", True):
-            return False
-        # MoE models need expert path, skip packed
-        if any(".experts." in name for name, _ in named_params_and_buffers(self.args, self.model)):
-            return False
-        # compressed-tensors needs pre/post process
-        if self.quantization_config and self.quantization_config.get("quant_method") == "compressed-tensors":
-            return False
-        return True
-
-    def _update_weights_vllm_packed(
-        self, converted_named_tensors: list[tuple[str, torch.Tensor]]
-    ) -> None:
-        """Single-shot vLLM weight update using packed broadcast (aligned with vLLM New Weight Syncing)."""
-        while not ray.get(self.rollout_engine_lock.acquire.remote()):
-            time.sleep(0.1)
-
-        try:
-            refs = update_weights_from_distributed(
-                self._group_name,
-                self._model_update_groups,
-                self.weight_version,
-                self.rollout_engines,
-                converted_named_tensors,
-                use_vllm=True,
-                packed=True,
-            )
-            ray.get(refs)
-        finally:
-            ray.get(self.rollout_engine_lock.release.remote())
-
     def _update_expert_weight_from_distributed(
         self,
         name: str,
@@ -455,7 +241,6 @@ def _update_bucket_weights_from_distributed(
             self.weight_version,
             self.rollout_engines,
             converted_named_tensors,
-            use_vllm=_is_vllm_backend(self.args),
         )
 
         ray.get(refs)
@@ -464,10 +249,6 @@ def _update_bucket_weights_from_distributed(
         pbar.update(1)
 
 
-def _is_vllm_backend(args: Namespace) -> bool:
-    return getattr(args, "rollout_backend", "sglang") == "vllm"
-
-
 def connect_rollout_engines_from_distributed(
     args: Namespace,
     group_name: str,
@@ -480,11 +261,6 @@ def connect_rollout_engines_from_distributed(
     ``engine_gpu_counts`` gives the number of GPUs per engine.  When engines
     have heterogeneous TP sizes (e.g. prefill TP=2, decode TP=4), each engine
     occupies a different number of ranks in the NCCL group.
-
-    For vLLM backend, the trainer-side NCCL communicator is created inside a
-    separate subprocess (_NcclBridge) to avoid conflicts between vLLM's raw
-    NCCL (PyNcclCommunicator) and the torch.distributed NCCL groups that
-    Megatron already holds in this process.
     """
     if engine_gpu_counts is None:
         engine_gpu_counts = [args.rollout_num_gpus_per_engine] * len(rollout_engines)
@@ -500,50 +276,24 @@ def connect_rollout_engines_from_distributed(
     for c in engine_gpu_counts:
         cumulative.append(cumulative[-1] + c)
 
-    # Fire engine init remotes first (non-blocking Ray calls).
     refs = [
         engine.init_weights_update_group.remote(
-            master_address=master_address,
-            master_port=master_port,
-            rank_offset=cumulative[i] + 1,
-            world_size=world_size,
-            group_name=group_name,
+            master_address,
+            master_port,
+            cumulative[i] + 1,
+            world_size,
+            group_name,
             backend="nccl",
         )
         for i, engine in enumerate(rollout_engines)
     ]
-
-    torch.cuda.synchronize()
-    torch.cuda.empty_cache()
-
-    if _is_vllm_backend(args):
-        # vLLM uses StatelessProcessGroup + PyNcclCommunicator (raw NCCL).
-        # Creating PyNcclCommunicator in the Megatron trainer process would
-        # conflict with torch.distributed's NCCL groups (issue #5477).
-        # Instead, we spawn a bridge subprocess that owns the PyNcclCommunicator,
-        # keeping the trainer process free of raw NCCL communicators.
-        device = torch.cuda.current_device()
-        logger.info(
-            "vLLM weight transfer via NcclBridge: addr=%s port=%d "
-            "world_size=%d device=%d CVD=%s",
-            master_address, master_port, world_size, device,
-            os.environ.get("CUDA_VISIBLE_DEVICES", "<unset>"),
-        )
-        model_update_groups = _NcclBridge(
-            master_address=master_address,
-            master_port=master_port,
-            world_size=world_size,
-            device=device,
-        )
-    else:
-        model_update_groups = init_process_group(
-            backend="nccl",
-            init_method=f"tcp://{master_address}:{master_port}",
-            world_size=world_size,
-            rank=0,
-            group_name=group_name,
-        )
-
+    model_update_groups = init_process_group(
+        backend="nccl",
+        init_method=f"tcp://{master_address}:{master_port}",
+        world_size=world_size,
+        rank=0,
+        group_name=group_name,
+    )
     ray.get(refs)
     return model_update_groups
 
@@ -553,56 +303,36 @@ def disconnect_rollout_engines_from_distributed(args, group_name, model_update_g
     Destroy NCCL on training and engines.
     """
     refs = [engine.destroy_weights_update_group.remote(group_name) for engine in rollout_engines]
-    if _is_vllm_backend(args):
-        if isinstance(model_update_groups, _NcclBridge):
-            model_update_groups.shutdown()
-        model_update_groups = None
-    else:
-        dist.destroy_process_group(model_update_groups)
+    dist.destroy_process_group(model_update_groups)
     ray.get(refs)
 
 
 def update_weights_from_distributed(
     group_name: str,
-    group,
+    group: dist.ProcessGroup,
     weight_version: int,
     rollout_engines: Sequence[ActorHandle],
     converted_named_tensors: Sequence[tuple[str, torch.Tensor]],
-    use_vllm: bool = False,
-    packed: bool = False,
 ) -> list[ObjectRef]:
     """
     Send metadata (Ray), broadcast tensors (NCCL rank 0 → engines).
-
-    For vLLM the *group* is an ``_NcclBridge`` instance (subprocess) so that
-    raw NCCL never runs inside the Megatron trainer process.
-    For sglang the *group* is a ``torch.distributed.ProcessGroup``.
     """
-    kwargs = {
-        "names": [name for name, _ in converted_named_tensors],
-        "dtypes": [param.dtype for _, param in converted_named_tensors],
-        "shapes": [param.shape for _, param in converted_named_tensors],
-        "group_name": group_name,
-        "weight_version": str(weight_version),
-    }
-    if use_vllm:
-        kwargs["packed"] = packed
-
     refs = [
-        engine.update_weights_from_distributed.remote(**kwargs)
+        engine.update_weights_from_distributed.remote(
+            names=[name for name, _ in converted_named_tensors],
+            dtypes=[param.dtype for _, param in converted_named_tensors],
+            shapes=[param.shape for _, param in converted_named_tensors],
+            group_name=group_name,
+            weight_version=str(weight_version),
+        )
         for engine in rollout_engines
     ]
 
-    if use_vllm and packed:
-        group.send_weights_packed(converted_named_tensors)
-    elif use_vllm:
-        group.broadcast_tensors([param.data for _, param in converted_named_tensors])
-    else:
-        handles = []
-        for _, param in converted_named_tensors:
-            handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
-        for handle in handles:
-            handle.wait()
+    handles = []
+    for _, param in converted_named_tensors:
+        handles.append(dist.broadcast(param.data, 0, group=group, async_op=True))
+    for handle in handles:
+        handle.wait()
 
     return refs
 
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
index dbba6aeb57..b53cc5bdc0 100644
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
+++ b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor.py
@@ -158,16 +158,9 @@ def update_weights(self) -> None:
         for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
             refs, long_lived_tensors = self._send_hf_params(hf_named_tensors)
             ray.get(refs)
-            # Free GPU tensors so the caching allocator can reuse the blocks,
-            # then release CUDA IPC cache entries whose consumers (sglang engines)
-            # have already closed their IPC handles.
-            del long_lived_tensors, hf_named_tensors
-            torch.cuda.ipc_collect()
+            del long_lived_tensors
 
         dist.barrier(group=get_gloo_group())
-        # After the barrier all engines have returned, so every rank's last-chunk
-        # IPC handles are now released by the consumers.  Clean them up.
-        torch.cuda.ipc_collect()
 
         # int4/fp4 post_process
         if rank == 0:
@@ -219,6 +212,7 @@ def _send_to_colocated_engine(
     if ipc_gather_group is None:
         return [], None
 
+    # TODO improve
     long_live_tensors = []
 
     if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
diff --git a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py b/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py
deleted file mode 100644
index d58d0377e7..0000000000
--- a/slime/backends/megatron_utils/update_weight/update_weight_from_tensor_vllm.py
+++ /dev/null
@@ -1,282 +0,0 @@
-"""
-UpdateVLLMWeightFromTensor
-==========================
-
-Update vLLM rollout engines using CUDA IPC (HTTP mode) when colocated on the
-same GPU(s) as the trainer.  This follows the vLLM RLHF HTTP IPC example:
-https://docs.vllm.ai/en/stable/examples/rl/rlhf_http_ipc/
-
-The flow:
-1. Megatron params → TP all-gather → HF conversion (via HfWeightIteratorBase)
-2. Send HF-named tensors to the vLLM server using
-   ``IPCWeightTransferEngine.trainer_send_weights()`` with
-   ``IPCTrainerSendWeightsArgs(mode="http", url=<vllm_url>)``.
-3. For any overflow (non-colocated) engines, fall back to NCCL distributed
-   broadcast identical to ``UpdateWeightFromDistributed``.
-
-The API is compatible with ``MegatronTrainRayActor`` — same ``__init__``,
-``connect_rollout_engines``, and ``update_weights`` signatures as
-``UpdateWeightFromTensor`` / ``UpdateWeightFromDistributed``.
-"""
-
-from __future__ import annotations
-
-import logging
-import os
-import time
-from argparse import Namespace
-from collections.abc import Callable, Mapping, Sequence
-
-import ray
-import torch
-import torch.distributed as dist
-from megatron.core import mpu
-from ray.actor import ActorHandle
-
-from slime.utils.distributed_utils import get_gloo_group
-
-from .common import all_gather_param, named_params_and_buffers
-from .hf_weight_iterator_base import HfWeightIteratorBase
-from .update_weight_from_distributed import (
-    connect_rollout_engines_from_distributed,
-    disconnect_rollout_engines_from_distributed,
-    post_process_weights,
-    update_weights_from_distributed,
-)
-
-logger = logging.getLogger(__name__)
-
-
-class UpdateVLLMWeightFromTensor:
-    """
-    Update colocated vLLM engines from tensor via CUDA IPC (HTTP mode).
-
-    Colocated path:
-        Megatron weights → TP all-gather → HF conversion → CUDA IPC to vLLM
-        server via ``IPCWeightTransferEngine.trainer_send_weights()``.
-
-    Non-colocated overflow path (optional):
-        Falls back to NCCL distributed broadcast via ``_NcclBridge``.
-    """
-
-    def __init__(
-        self,
-        args: Namespace,
-        model: Sequence[torch.nn.Module],
-        weights_getter: Callable[[], Mapping[str, torch.Tensor]],
-        *,
-        model_name: str,
-        quantization_config: dict[str, int | str | list[str]] | None,
-    ) -> None:
-        self.args = args
-        self.model = model
-        self.weights_getter = weights_getter
-        self.model_name = model_name
-        self.quantization_config = quantization_config
-        self.weight_version = 0
-
-        self._hf_weight_iterator = HfWeightIteratorBase.create(
-            args=args, model=model, model_name=model_name, quantization_config=quantization_config,
-        )
-
-        # Populated by connect_rollout_engines
-        self.rollout_engines: list[ActorHandle] = []
-        self._colocated_engines: list[ActorHandle] = []
-        self._colocated_vllm_urls: list[str] = []
-        self._distributed_engines: list[ActorHandle] = []
-        self._model_update_groups = None
-        self._is_distributed_src_rank = False
-        self._group_name = "slime"
-        self._ipc_initialized = False
-
-    # ------------------------------------------------------------------
-    # connect / disconnect
-    # ------------------------------------------------------------------
-
-    def connect_rollout_engines(
-        self,
-        rollout_engines: Sequence[ActorHandle],
-        rollout_engine_lock: ActorHandle,
-        engine_gpu_counts: Sequence[int] | None = None,
-        engine_gpu_offsets: Sequence[int] | None = None,
-    ) -> None:
-        """
-        Split colocated / distributed engines.
-
-        For colocated engines we resolve their vLLM base URLs (needed by the
-        IPC weight transfer HTTP mode).  For overflow distributed engines we
-        create the NCCL bridge just like ``UpdateWeightFromDistributed``.
-        """
-        self.rollout_engines = list(rollout_engines)
-        self.rollout_engine_lock = rollout_engine_lock
-
-        if engine_gpu_counts is None:
-            engine_gpu_counts = [self.args.rollout_num_gpus_per_engine] * len(rollout_engines)
-        if engine_gpu_offsets is None:
-            engine_gpu_offsets = []
-            offset = 0
-            for c in engine_gpu_counts:
-                engine_gpu_offsets.append(offset)
-                offset += c
-
-        # Determine colocated vs distributed engines
-        total_actor_gpus = self.args.actor_num_nodes * self.args.actor_num_gpus_per_node
-        colocate_engine_nums = 0
-        for gpu_offset, gpu_count in zip(engine_gpu_offsets, engine_gpu_counts, strict=True):
-            if gpu_offset + gpu_count > total_actor_gpus:
-                break
-            colocate_engine_nums += 1
-
-        self._colocated_engines = list(rollout_engines[:colocate_engine_nums])
-        self._distributed_engines = list(rollout_engines[colocate_engine_nums:])
-
-        # Resolve vLLM base URLs for colocated engines (blocking Ray call)
-        if dist.get_rank() == 0 and self._colocated_engines:
-            url_refs = [engine.get_vllm_url.remote() for engine in self._colocated_engines]
-            self._colocated_vllm_urls = ray.get(url_refs)
-            logger.info("Colocated vLLM URLs for IPC weight transfer: %s", self._colocated_vllm_urls)
-
-        # Set up NCCL bridge for distributed overflow engines
-        if self._distributed_engines:
-            distributed_gpu_counts = engine_gpu_counts[colocate_engine_nums:]
-            self._is_distributed_src_rank = (
-                mpu.get_data_parallel_rank(with_context_parallel=True) == 0
-                and mpu.get_tensor_model_parallel_rank() == 0
-                and mpu.get_pipeline_model_parallel_rank() == 0
-            )
-            if self._is_distributed_src_rank:
-                if self._model_update_groups is not None:
-                    disconnect_rollout_engines_from_distributed(
-                        self.args, self._group_name, self._model_update_groups, self._distributed_engines,
-                    )
-                self._model_update_groups = connect_rollout_engines_from_distributed(
-                    self.args,
-                    self._group_name,
-                    self._distributed_engines,
-                    engine_gpu_counts=distributed_gpu_counts,
-                )
-
-    # ------------------------------------------------------------------
-    # weight update
-    # ------------------------------------------------------------------
-
-    @torch.no_grad()
-    def update_weights(self) -> None:
-        """
-        Main entry-point called by ``MegatronTrainRayActor.update_weights()``.
-
-        Pause → flush → init IPC (once) → send weights via IPC → resume.
-        """
-
-        self.weight_version += 1
-        rank = dist.get_rank()
-
-        # Pause generation and flush KV cache on all engines
-        if rank == 0:
-            all_engines = self._colocated_engines + self._distributed_engines
-            ray.get([engine.pause_generation.remote() for engine in all_engines])
-            ray.get([engine.flush_cache.remote() for engine in all_engines])
-            if self.quantization_config and self.quantization_config.get("quant_method") in ["compressed-tensors"]:
-                post_process_weights(
-                    restore_weights_before_load=True,
-                    post_process_quantization=False,
-                    rollout_engines=all_engines,
-                )
-        dist.barrier(group=get_gloo_group())
-
-        # Initialize IPC weight transfer engines (first time only)
-        if rank == 0 and self._colocated_engines and not self._ipc_initialized:
-            self._init_ipc_weight_transfer()
-            self._ipc_initialized = True
-
-        # Get megatron weights and iterate HF chunks
-        megatron_local_weights = self.weights_getter()
-
-        for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
-            # Send to colocated engines via IPC
-            if rank == 0 and self._colocated_engines:
-                self._send_via_ipc(hf_named_tensors)
-
-            # Send to distributed engines via NCCL
-            if self._distributed_engines and self._is_distributed_src_rank:
-                refs = update_weights_from_distributed(
-                    self._group_name,
-                    self._model_update_groups,
-                    self.weight_version,
-                    self._distributed_engines,
-                    hf_named_tensors,
-                    use_vllm=True,
-                    packed=True,
-                )
-                if refs:
-                    ray.get(refs)
-
-        dist.barrier(group=get_gloo_group())
-
-        # Post-process and resume
-        if rank == 0:
-            all_engines = self._colocated_engines + self._distributed_engines
-            if self.quantization_config and self.quantization_config.get("quant_method") in ["compressed-tensors"]:
-                post_process_weights(
-                    restore_weights_before_load=False,
-                    post_process_quantization=True,
-                    rollout_engines=all_engines,
-                )
-            # Bump weight version on sidecar for colocated engines
-            for engine in self._colocated_engines:
-                try:
-                    ray.get(engine.set_weight_version.remote(self.weight_version))
-                except Exception as exc:
-                    logger.warning("Failed to set weight version on engine: %s", exc)
-
-            ray.get([engine.continue_generation.remote() for engine in all_engines])
-        dist.barrier(group=get_gloo_group())
-
-    # ------------------------------------------------------------------
-    # IPC helpers
-    # ------------------------------------------------------------------
-
-    def _init_ipc_weight_transfer(self) -> None:
-        """
-        Call ``/init_weight_transfer_engine`` on each colocated vLLM server.
-        For IPC backend this is a no-op on the server side but is still required
-        by vLLM's weight transfer protocol.
-        """
-        import requests
-
-        for url in self._colocated_vllm_urls:
-            try:
-                resp = requests.post(
-                    f"{url}/init_weight_transfer_engine",
-                    json={"init_info": {}},
-                    timeout=60,
-                )
-                resp.raise_for_status()
-                logger.info("Initialized IPC weight transfer on %s", url)
-            except Exception as exc:
-                logger.error("Failed to init IPC weight transfer on %s: %s", url, exc)
-                raise
-
-    def _send_via_ipc(self, hf_named_tensors: list[tuple[str, torch.Tensor]]) -> None:
-        """
-        Send HF-named tensors to all colocated vLLM engines via CUDA IPC.
-
-        Uses ``vllm.distributed.weight_transfer.ipc_engine.IPCWeightTransferEngine``
-        with ``mode="http"`` so the IPC handles are sent via HTTP to the vLLM
-        server's ``/update_weights`` endpoint.
-        """
-        # Allow insecure serialization for IPC handle serialization
-        os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
-
-        from vllm.distributed.weight_transfer.ipc_engine import (
-            IPCTrainerSendWeightsArgs,
-            IPCWeightTransferEngine,
-        )
-
-        for url in self._colocated_vllm_urls:
-            trainer_args = IPCTrainerSendWeightsArgs(mode="http", url=url)
-            IPCWeightTransferEngine.trainer_send_weights(
-                iterator=iter(hf_named_tensors),
-                trainer_args=trainer_args,
-            )
-            logger.debug("IPC weight transfer completed to %s", url)
diff --git a/slime/backends/megatron_utils/weight_sync_utils.py b/slime/backends/megatron_utils/weight_sync_utils.py
deleted file mode 100644
index 1626444ad8..0000000000
--- a/slime/backends/megatron_utils/weight_sync_utils.py
+++ /dev/null
@@ -1,312 +0,0 @@
-"""
-Local reimplementations of sglang weight-sync utilities.
-
-Used as fallback when sglang is not installed (e.g. vLLM-only mode).
-The three classes/functions here are API-compatible with their sglang
-counterparts so that the rest of the megatron weight-update code can
-work unchanged.
-
-Origin (sglang):
-  - FlattenedTensorBucket  → sglang.srt.weight_sync.tensor_bucket
-  - MultiprocessingSerializer / SafeUnpickler → sglang.srt.utils.common
-  - monkey_patch_torch_reductions → sglang.srt.utils.patch_torch
-"""
-
-from __future__ import annotations
-
-import base64
-import io
-import pickle
-from dataclasses import dataclass
-from multiprocessing.reduction import ForkingPickler
-from typing import Callable, Union
-
-import torch
-from torch.multiprocessing import reductions
-
-# ── FlattenedTensorBucket ───────────────────────────────────────────
-
-
-@dataclass
-class FlattenedTensorMetadata:
-    """Metadata for a tensor in a flattened bucket."""
-
-    name: str
-    shape: torch.Size
-    dtype: torch.dtype
-    start_idx: int
-    end_idx: int
-    numel: int
-
-
-class FlattenedTensorBucket:
-    """
-    A bucket that flattens multiple tensors into a single uint8 tensor
-    for efficient serialisation, while preserving all metadata needed
-    for reconstruction.
-
-    API-compatible with ``sglang.srt.weight_sync.tensor_bucket.FlattenedTensorBucket``.
-    """
-
-    # Checked by callers to decide whether to group tensors by dtype.
-    supports_multi_dtypes = True
-
-    def __init__(
-        self,
-        named_tensors: list[tuple[str, torch.Tensor]] | None = None,
-        flattened_tensor: torch.Tensor | None = None,
-        metadata: list[FlattenedTensorMetadata] | None = None,
-    ):
-        if named_tensors is not None:
-            if not named_tensors:
-                raise ValueError("Cannot create empty tensor bucket")
-
-            self.metadata: list[FlattenedTensorMetadata] = [None] * len(named_tensors)
-            current_idx = 0
-            flat_parts: list[torch.Tensor] = [None] * len(named_tensors)
-
-            for i, (name, tensor) in enumerate(named_tensors):
-                flat = tensor.flatten().view(torch.uint8)
-                numel = flat.numel()
-                flat_parts[i] = flat
-                self.metadata[i] = FlattenedTensorMetadata(
-                    name=name,
-                    shape=tensor.shape,
-                    dtype=tensor.dtype,
-                    start_idx=current_idx,
-                    end_idx=current_idx + numel,
-                    numel=numel,
-                )
-                current_idx += numel
-
-            self.flattened_tensor: torch.Tensor = torch.cat(flat_parts, dim=0)
-        else:
-            if flattened_tensor is None or metadata is None:
-                raise ValueError(
-                    "Must provide either named_tensors or both flattened_tensor and metadata"
-                )
-            self.flattened_tensor = flattened_tensor
-            self.metadata = metadata
-
-    def get_flattened_tensor(self) -> torch.Tensor:
-        """Return the single flat uint8 tensor."""
-        return self.flattened_tensor
-
-    def get_metadata(self) -> list[FlattenedTensorMetadata]:
-        """Return per-tensor metadata list."""
-        return self.metadata
-
-    def reconstruct_tensors(self) -> list[tuple[str, torch.Tensor]]:
-        """Reconstruct the original named tensors from the flat representation."""
-        reconstructed = [None] * len(self.metadata)
-        for i, meta in enumerate(self.metadata):
-            tensor = (
-                self.flattened_tensor[meta.start_idx : meta.end_idx]
-                .view(meta.dtype)
-                .reshape(meta.shape)
-            )
-            reconstructed[i] = (meta.name, tensor)
-        return reconstructed
-
-
-# ── SafeUnpickler / MultiprocessingSerializer ───────────────────────
-
-
-class SafeUnpickler(pickle.Unpickler):
-    """
-    Unpickler with an allow-list to prevent arbitrary code execution.
-
-    API-compatible with the ``SafeUnpickler`` in ``sglang.srt.utils.common``.
-    """
-
-    ALLOWED_MODULE_PREFIXES = {
-        # Python builtins (specific safe classes only – see ALLOW_CLASSES)
-        "collections.",
-        "copyreg.",
-        "functools.",
-        "itertools.",
-        "operator.",
-        "types.",
-        "weakref.",
-        # PyTorch
-        "torch.",
-        "torch._tensor.",
-        "torch.storage.",
-        "torch.nn.parameter.",
-        "torch.autograd.function.",
-        # torch.distributed
-        "torch.distributed.",
-        "torch.distributed._shard.",
-        "torch.distributed._composable.",
-        "torch._C._distributed_c10d.",
-        "torch._C._distributed_fsdp.",
-        "torch.distributed.optim.",
-        # multiprocessing
-        "multiprocessing.resource_sharer.",
-        "multiprocessing.reduction.",
-        "pickletools.",
-        # HuggingFace / PEFT
-        "peft.",
-        "transformers.",
-        "huggingface_hub.",
-        # slime local reimplementation
-        "slime.backends.megatron_utils.weight_sync_utils.",
-        # sglang (if installed alongside)
-        "sglang.srt.weight_sync.tensor_bucket.",
-        "sglang.srt.model_executor.model_runner.",
-        "sglang.srt.layers.",
-        "sglang.srt.utils.",
-        # NPU
-        "torch_npu.",
-    }
-
-    # Specific builtins classes that are safe to unpickle.
-    ALLOW_CLASSES = {
-        ("builtins", "True"),
-        ("builtins", "False"),
-        ("builtins", "None"),
-        ("builtins", "dict"),
-        ("builtins", "list"),
-        ("builtins", "tuple"),
-        ("builtins", "set"),
-        ("builtins", "frozenset"),
-        ("builtins", "int"),
-        ("builtins", "float"),
-        ("builtins", "complex"),
-        ("builtins", "str"),
-        ("builtins", "bytes"),
-        ("builtins", "bytearray"),
-        ("builtins", "bool"),
-        ("builtins", "slice"),
-        ("builtins", "range"),
-        ("builtins", "enumerate"),
-        ("builtins", "map"),
-        ("builtins", "zip"),
-        ("builtins", "filter"),
-        ("builtins", "reversed"),
-        ("builtins", "sorted"),
-    }
-
-    DENY_CLASSES = {
-        ("builtins", "eval"),
-        ("builtins", "exec"),
-        ("builtins", "compile"),
-        ("builtins", "getattr"),
-        ("builtins", "setattr"),
-        ("builtins", "delattr"),
-        ("builtins", "__import__"),
-        ("builtins", "globals"),
-        ("builtins", "locals"),
-        ("builtins", "open"),
-        ("builtins", "breakpoint"),
-        ("builtins", "input"),
-        ("builtins", "memoryview"),
-        ("os", "system"),
-        ("subprocess", "Popen"),
-        ("subprocess", "run"),
-        ("codecs", "decode"),
-        ("types", "CodeType"),
-        ("types", "FunctionType"),
-    }
-
-    def find_class(self, module: str, name: str):
-        if (module, name) in self.DENY_CLASSES:
-            raise RuntimeError(
-                f"Blocked unsafe class loading ({module}.{name}), "
-                f"to prevent exploitation of CVE-2025-10164"
-            )
-        # Check explicit allow-list for builtins (strict whitelist)
-        if module == "builtins":
-            if (module, name) in self.ALLOW_CLASSES:
-                return super().find_class(module, name)
-            raise RuntimeError(
-                f"Blocked unsafe class loading ({module}.{name}), "
-                f"to prevent exploitation of CVE-2025-10164"
-            )
-        if any((module + ".").startswith(prefix) for prefix in self.ALLOWED_MODULE_PREFIXES):
-            return super().find_class(module, name)
-        raise RuntimeError(
-            f"Blocked unsafe class loading ({module}.{name}), "
-            f"to prevent exploitation of CVE-2025-10164"
-        )
-
-
-class MultiprocessingSerializer:
-    """
-    Serialize / deserialize Python objects via ``ForkingPickler`` so that
-    CUDA tensors are transferred through shared memory (IPC handles).
-
-    API-compatible with ``sglang.srt.utils.common.MultiprocessingSerializer``.
-
-    Uses stdlib ``base64`` instead of ``pybase64`` to avoid adding a dependency.
-    """
-
-    @staticmethod
-    def serialize(obj, output_str: bool = False):
-        buf = io.BytesIO()
-        ForkingPickler(buf).dump(obj)
-        buf.seek(0)
-        output = buf.read()
-        if output_str:
-            output = base64.b64encode(output).decode("utf-8")
-        return output
-
-    @staticmethod
-    def deserialize(data):
-        if isinstance(data, str):
-            data = base64.b64decode(data, validate=True)
-        return SafeUnpickler(io.BytesIO(data)).load()
-
-
-# ── monkey_patch_torch_reductions ───────────────────────────────────
-
-_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6
-
-
-def _device_to_uuid(device: int) -> str:
-    return str(torch.cuda.get_device_properties(device).uuid)
-
-
-def _device_from_maybe_uuid(device_maybe_uuid: Union[int, str]) -> int:
-    if isinstance(device_maybe_uuid, int):
-        return device_maybe_uuid
-    if isinstance(device_maybe_uuid, str):
-        for device in range(torch.cuda.device_count()):
-            if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid:
-                return device
-        raise RuntimeError("Invalid device_uuid=" + device_maybe_uuid)
-    raise RuntimeError(f"Unknown type: {device_maybe_uuid=}")
-
-
-def _modify_tuple(t, index: int, modifier: Callable):
-    return (*t[:index], modifier(t[index]), *t[index + 1 :])
-
-
-def _reduce_tensor_modified(*args, **kwargs):
-    output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs)
-    output_args = _modify_tuple(output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid)
-    return output_fn, output_args
-
-
-def _rebuild_cuda_tensor_modified(*args):
-    args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid)
-    return reductions._rebuild_cuda_tensor_original(*args)
-
-
-def monkey_patch_torch_reductions():
-    """
-    Monkey-patch ``torch.multiprocessing.reductions`` so that CUDA tensors
-    are identified by device UUID rather than ordinal index.
-
-    This works around https://github.com/pytorch/pytorch/pull/149248.
-
-    API-compatible with ``sglang.srt.utils.patch_torch.monkey_patch_torch_reductions``.
-    """
-    if hasattr(reductions, "_reduce_tensor_original"):
-        return  # already patched
-    reductions._reduce_tensor_original = reductions.reduce_tensor
-    reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
-
-    reductions.reduce_tensor = _reduce_tensor_modified
-    reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
-    reductions.init_reductions()
diff --git a/slime/backends/sglang_utils/arguments.py b/slime/backends/sglang_utils/arguments.py
index c955583b15..71543e02d3 100644
--- a/slime/backends/sglang_utils/arguments.py
+++ b/slime/backends/sglang_utils/arguments.py
@@ -21,6 +21,12 @@ def add_sglang_router_arguments(parser):
         default=None,
         help="Port of the SGLang router",
     )
+    parser.add_argument(
+        "--sglang-router-policy",
+        type=str,
+        default=None,
+        help="Routing policy for the SGLang router (e.g., 'consistent_hashing', 'round_robin')",
+    )
     parser.add_argument(
         "--sglang-router-request-timeout-secs",
         type=int,
@@ -36,6 +42,7 @@ def add_sglang_arguments(parser):
     """
     parser = add_sglang_router_arguments(parser)
     parser.set_defaults(router_balance_abs_threshold=10, router_balance_rel_threshold=1.2)
+    parser.add_argument("--sglang-server-concurrency", type=int, default=512)
 
     old_add_argument = parser.add_argument
 
@@ -124,7 +131,7 @@ def new_add_argument_wrapper(*name_or_flags, **kwargs):
         default=None,
         help=(
             "Path to a YAML config for SGLang engine deployment. "
-            "Defines server_groups with worker_type (regular/prefill/decode/placeholder), "
+            "Defines engine_groups with worker_type (regular/prefill/decode/placeholder), "
             "num_gpus per group, and optional per-group 'overrides' dict of "
             "ServerArgs field names that override the base --sglang-* CLI args. "
             "Placeholder groups reserve GPU slots without creating engines. "
@@ -167,14 +174,14 @@ def validate_args(args):
 
     assert not (
         getattr(args, "sglang_config", None) is not None and getattr(args, "prefill_num_servers", None) is not None
-    ), "sglang_config and prefill_num_servers are mutually exclusive. Use server_groups in the YAML config instead."
+    ), "sglang_config and prefill_num_servers are mutually exclusive. Use engine_groups in the YAML config instead."
 
 
 def sglang_parse_args():
     """
     Parse sglang server arguments independently using a separate ArgumentParser.
     Uses parse_known_args() to only consume sglang-related arguments from sys.argv,
-    allowing the remaining arguments to be parsed by megatron separately.
+    allowing the remaining arguments to be parsed by megatron/fsdp separately.
 
     Returns:
         argparse.Namespace: Parsed sglang arguments (all attributes prefixed with sglang_).
diff --git a/slime/backends/sglang_utils/sglang_config.py b/slime/backends/sglang_utils/sglang_config.py
deleted file mode 100644
index 827b47a1bd..0000000000
--- a/slime/backends/sglang_utils/sglang_config.py
+++ /dev/null
@@ -1,207 +0,0 @@
-"""Configuration dataclasses for SGLang engine deployment."""
-
-import dataclasses
-import logging
-
-import yaml
-
-logger = logging.getLogger(__name__)
-
-
-@dataclasses.dataclass
-class ServerGroupConfig:
-    """Configuration for a single server group.
-
-    Attributes:
-        worker_type: One of "regular", "prefill", "decode", "placeholder",
-                     or "encoder".
-                     "placeholder" reserves GPU slots without creating engines.
-                     "encoder" creates encoder-only engines for EPD
-                     (Encoder-Prefill-Decode) disaggregation; encoder engines
-                     are started first and their URLs are automatically
-                     injected into prefill groups as ``encoder_urls``.
-        num_gpus: Total number of GPUs for this group.
-        num_gpus_per_engine: GPUs per engine for this group.  Overrides the
-                             model-level or global ``--rollout-num-gpus-per-engine``.
-        overrides: Optional dict of SGLang ``ServerArgs`` field overrides.
-                   These are applied on top of the base CLI ``--sglang-*``
-                   arguments in ``_compute_server_args``.
-    """
-
-    worker_type: str
-    num_gpus: int
-    num_gpus_per_engine: int | None = None
-    overrides: dict = dataclasses.field(default_factory=dict)
-
-    def __post_init__(self):
-        valid_types = {"regular", "prefill", "decode", "placeholder", "encoder"}
-        assert (
-            self.worker_type in valid_types
-        ), f"Invalid worker_type '{self.worker_type}', must be one of {valid_types}"
-        assert self.num_gpus > 0, f"num_gpus must be > 0, got {self.num_gpus}"
-
-
-@dataclasses.dataclass
-class ModelConfig:
-    """Configuration for a single model deployment.
-
-    Attributes:
-        name: Unique name for this model (e.g. "actor", "reward").
-        model_path: HF checkpoint path.  Falls back to ``args.hf_checkpoint``.
-        num_gpus_per_engine: Default GPUs per engine for all groups in this
-                             model.  Individual groups can override.
-        server_groups: Server group configurations for this model.
-        update_weights: Whether this model receives weight updates from
-                        training.  Set to ``False`` for frozen models
-                        (reference, reward, etc.).  When ``None`` (default),
-                        automatically inferred in ``resolve()``: ``True`` if
-                        model_path matches ``args.hf_checkpoint``, ``False``
-                        otherwise.
-    """
-
-    name: str
-    model_path: str | None = None
-    num_gpus_per_engine: int | None = None
-    server_groups: list[ServerGroupConfig] = dataclasses.field(default_factory=list)
-    update_weights: bool | None = None
-
-    def resolve(self, args) -> None:
-        """Resolve per-group defaults from model-level then args-level values."""
-        default_gpus_per_engine = self.num_gpus_per_engine or args.rollout_num_gpus_per_engine
-        default_model_path = self.model_path or args.hf_checkpoint
-        for g in self.server_groups:
-            if g.num_gpus_per_engine is None:
-                g.num_gpus_per_engine = default_gpus_per_engine
-            # Inject model_path into overrides so _compute_server_args picks it up.
-            if "model_path" not in g.overrides:
-                g.overrides["model_path"] = default_model_path
-
-        # Validate: all server groups within a model must share the same model_path.
-        if self.server_groups:
-            model_paths = {g.overrides["model_path"] for g in self.server_groups}
-            assert len(model_paths) == 1, (
-                f"Model '{self.name}' has server groups with different model_path values: "
-                f"{model_paths}. All server groups within a model must use the same model_path."
-            )
-            effective_model_path = model_paths.pop()
-        else:
-            effective_model_path = default_model_path
-
-        # Auto-infer update_weights when not explicitly set.
-        if self.update_weights is None:
-            if effective_model_path != args.hf_checkpoint:
-                logger.warning(
-                    f"Model '{self.name}' uses model_path='{effective_model_path}' which differs "
-                    f"from hf_checkpoint='{args.hf_checkpoint}'. Defaulting update_weights to False. "
-                    f"Set update_weights explicitly in the config to suppress this warning."
-                )
-                self.update_weights = False
-            else:
-                self.update_weights = True
-
-    @property
-    def has_pd_disaggregation(self) -> bool:
-        return any(g.worker_type in ("prefill", "decode") for g in self.server_groups)
-
-    @property
-    def has_encoder_disaggregation(self) -> bool:
-        return any(g.worker_type == "encoder" for g in self.server_groups)
-
-    @property
-    def total_num_gpus(self) -> int:
-        return sum(g.num_gpus for g in self.server_groups)
-
-
-@dataclasses.dataclass
-class SglangConfig:
-    """Configuration for SGLang engine deployment.
-
-    Loaded from ``--sglang-config`` YAML file.
-
-    **Config format**::
-
-        sglang:
-          - name: actor
-            model_path: /path/to/actor
-            update_weights: true          # receives training weight updates (default)
-            num_gpus_per_engine: 2
-            server_groups:
-              - worker_type: prefill
-                num_gpus: 4
-                num_gpus_per_engine: 2
-              - worker_type: decode
-                num_gpus: 8
-                num_gpus_per_engine: 4
-          - name: ref
-            model_path: /path/to/ref
-            update_weights: false          # frozen, no weight updates
-            server_groups:
-              - worker_type: regular
-                num_gpus: 4
-
-    Each model gets its own router.  ``placeholder`` groups reserve GPU
-    slots without creating engines.  ``overrides`` are ``ServerArgs``
-    field names applied on top of the base ``--sglang-*`` CLI args.
-
-    Set ``update_weights: false`` for frozen models (reference, reward,
-    etc.) that should not receive weight updates from training.
-
-    .. note::
-
-       ``engine_groups`` is accepted as a backward-compatible alias for
-       ``server_groups`` in the YAML config.
-    """
-
-    models: list[ModelConfig]
-
-    @staticmethod
-    def from_yaml(path: str) -> "SglangConfig":
-        with open(path) as f:
-            data = yaml.safe_load(f)
-
-        assert "sglang" in data, (
-            f"sglang config must have a 'sglang' key, got {list(data.keys())}. "
-            f"Wrap your server_groups inside a model entry under 'sglang'."
-        )
-        models = []
-        for m in data["sglang"]:
-            # Accept both "server_groups" and legacy "engine_groups".
-            raw_groups = m.get("server_groups") or m.get("engine_groups") or []
-            groups = [ServerGroupConfig(**g) for g in raw_groups]
-            models.append(
-                ModelConfig(
-                    name=m["name"],
-                    model_path=m.get("model_path"),
-                    num_gpus_per_engine=m.get("num_gpus_per_engine"),
-                    server_groups=groups,
-                    update_weights=m.get("update_weights"),
-                )
-            )
-        return SglangConfig(models=models)
-
-    @staticmethod
-    def from_prefill_num_servers(args) -> "SglangConfig":
-        """Build a config equivalent to the legacy --prefill-num-servers flag."""
-        total_gpus = args.rollout_num_gpus
-        prefill_gpus = args.prefill_num_servers * args.rollout_num_gpus_per_engine
-        decode_gpus = total_gpus - prefill_gpus
-        assert decode_gpus > 0, f"No decode GPUs: total {total_gpus}, prefill {prefill_gpus}"
-        return SglangConfig(
-            models=[
-                ModelConfig(
-                    name="default",
-                    server_groups=[
-                        ServerGroupConfig(worker_type="prefill", num_gpus=prefill_gpus),
-                        ServerGroupConfig(worker_type="decode", num_gpus=decode_gpus),
-                    ],
-                )
-            ]
-        )
-
-    @property
-    def has_pd_disaggregation(self) -> bool:
-        return any(m.has_pd_disaggregation for m in self.models)
-
-    @property
-    def total_num_gpus(self) -> int:
-        return sum(m.total_num_gpus for m in self.models)
diff --git a/slime/backends/sglang_utils/sglang_engine.py b/slime/backends/sglang_utils/sglang_engine.py
index c28e13d5f9..290b1fcdce 100644
--- a/slime/backends/sglang_utils/sglang_engine.py
+++ b/slime/backends/sglang_utils/sglang_engine.py
@@ -26,6 +26,9 @@ def get_base_gpu_id(args, rank):
     else:
         num_actor_gpus = 0 if args.debug_rollout_only else args.actor_num_gpus_per_node * args.actor_num_nodes
         start_index = (num_actor_gpus + rank * num_gpus) % args.num_gpus_per_node
+        if args.use_critic:
+            num_critic_gpus = args.critic_num_gpus_per_node * args.critic_num_nodes
+            start_index = (num_actor_gpus + num_critic_gpus + rank * num_gpus) % args.num_gpus_per_node
     return start_index
 
 
@@ -48,15 +51,6 @@ def _to_local_gpu_id(physical_gpu_id: int) -> int:
 
 
 def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
-    if getattr(server_args, "encoder_only", False):
-        from sglang.srt.disaggregation.encode_server import launch_server_process as sglang_launch_server_process
-
-        return sglang_launch_server_process(
-            server_args,
-            start_method="spawn",
-            wait_for_server=True,
-        )
-
     from sglang.srt.entrypoints.http_server import launch_server
 
     multiprocessing.set_start_method("spawn", force=True)
@@ -64,8 +58,8 @@ def launch_server_process(server_args: ServerArgs) -> multiprocessing.Process:
     p = multiprocessing.Process(target=launch_server, args=(server_args,))
     p.start()
 
-    if getattr(server_args, "node_rank", 0) != 0:
-        return p
+    if server_args.node_rank != 0:
+        return
 
     _wait_server_healthy(
         base_url=server_args.url(),
@@ -96,6 +90,21 @@ def _wait_server_healthy(base_url, api_key, is_process_alive):
 
             time.sleep(2)
 
+        # use flush_cache to make sure the working queue is empty, so that we can do offload
+        while True:
+            try:
+                response = session.get(f"{base_url}/flush_cache", headers=headers)
+                if response.status_code == 200:
+                    break
+
+            except requests.RequestException:
+                pass
+
+            if not is_process_alive():
+                raise Exception("Server process terminated unexpectedly.")
+
+            time.sleep(2)
+
 
 class SGLangEngine(RayActor):
     def __init__(
@@ -194,14 +203,13 @@ def _init_normal(self, server_args_dict):
         logger.info(f"Launch HttpServerEngineAdapter at: {self.server_host}:{self.server_port}")
         self.process = launch_server_process(ServerArgs(**server_args_dict))
 
-        if self.worker_type == "encoder":
-            return
-
         if self.node_rank == 0 and self.router_ip and self.router_port:
-            if parse(sglang_router.__version__) <= parse("0.2.1"):
-                assert self.worker_type == "regular", "pd disaggregation is not supported in old router."
+            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
+                assert (
+                    self.worker_type == "regular"
+                ), "pd disaggregation is not supported in old router or slime router."
                 response = requests.post(
-                    f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}",
+                    f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}"
                 )
             else:
                 payload = {
@@ -304,20 +312,15 @@ def flush_cache(self):
         else:
             raise TimeoutError("Timeout while flushing cache.")
 
-    def get_url(self):
-        if self.node_rank != 0:
-            return None
-        return f"http://{self.server_host}:{self.server_port}"
-
     def shutdown(self):
         if self.args.rollout_external:
             return
 
         logger.info(f"Shutdown engine {self.server_host}:{self.server_port}...")
-        if self.worker_type != "encoder" and self.node_rank == 0:
+        if self.node_rank == 0:
             worker_url = f"http://{self.server_host}:{self.server_port}"
             response = None
-            if parse(sglang_router.__version__) <= parse("0.2.1"):
+            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
                 response = requests.post(
                     f"http://{self.router_ip}:{self.router_port}/remove_worker?url=http://{self.server_host}:{self.server_port}"
                 )
@@ -367,17 +370,6 @@ def resume_memory_occupation(self, tags: list[str] = None):
     def check_weights(self, action: str):
         return self._make_request("weights_checker", {"action": action})
 
-    def update_weights_from_disk(self, model_path: str, load_format: str | None = None):
-        """Reload weights from *model_path* without restarting the engine.
-
-        Used for non-updatable (frozen) models that overlap with megatron:
-        after offload, weights are restored from disk instead of CPU cache.
-        """
-        payload = {"model_path": model_path}
-        if load_format is not None:
-            payload["load_format"] = load_format
-        return self._make_request("update_weights_from_disk", payload)
-
     def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name, backend):
         return self._make_request(
             "init_weights_update_group",
@@ -537,14 +529,11 @@ def _compute_server_args(
         "skip_server_warmup": True,
         # always enable draft weights cpu backup so that we run training without mtp weights.
         "enable_draft_weights_cpu_backup": True,
-        # Always enable Prometheus metrics so the /engine_metrics endpoint is
-        # available for W&B scraping (regardless of --sglang-enable-metrics).
-        "enable_metrics": True,
     }
 
     if worker_type == "prefill":
         kwargs["disaggregation_mode"] = "prefill"
-        kwargs["load_balance_method"] = "follow_bootstrap_room"
+        kwargs["load_balance_method"] = "round_robin"
         assert (
             disaggregation_bootstrap_port is not None
         ), "disaggregation_bootstrap_port must be set for prefill worker"
@@ -552,8 +541,6 @@ def _compute_server_args(
     elif worker_type == "decode":
         kwargs["disaggregation_mode"] = "decode"
         kwargs["prefill_round_robin_balance"] = True
-    elif worker_type == "encoder":
-        kwargs["encoder_only"] = True
 
     if args.use_rollout_routing_replay:
         kwargs["enable_return_routed_experts"] = True
@@ -561,35 +548,22 @@ def _compute_server_args(
         kwargs["dtype"] = "float16"
     external_engine_need_check_fields = [k for k in kwargs.keys() if k not in _EXTERNAL_ENGINE_SKIP_CHECK_FIELDS]
 
-    server_arg_fields = dataclasses.fields(ServerArgs)
-    server_arg_field_names = {attr.name for attr in server_arg_fields}
     unused_keys = set(kwargs.keys())
-    for attr in server_arg_fields:
+    for attr in dataclasses.fields(ServerArgs):
         if worker_type == "decode" and attr.name == "enable_hierarchical_cache":
             continue
         if hasattr(args, f"sglang_{attr.name}") and attr.name not in kwargs:
             kwargs[attr.name] = getattr(args, f"sglang_{attr.name}")
         unused_keys.discard(attr.name)
 
-    # Per-server-group overrides from --sglang-config YAML.
+    # Per-engine-group overrides from --sglang-config YAML.
     # Applied after base args so they take highest priority.
     if sglang_overrides:
         for key, value in sglang_overrides.items():
-            normalized_key = key.replace("-", "_")
-            if normalized_key != key:
-                logger.warning(
-                    f"sglang_overrides key '{key}' normalized to '{normalized_key}' (rank={rank}). "
-                    "Please use underscore style in YAML overrides."
-                )
-            if normalized_key in kwargs:
-                logger.info(
-                    f"sglang_overrides: overriding {normalized_key}={kwargs[normalized_key]} -> {value} (rank={rank})"
-                )
-            kwargs[normalized_key] = value
-            if normalized_key in server_arg_field_names:
-                unused_keys.discard(normalized_key)
-            else:
-                unused_keys.add(normalized_key)
+            if key in kwargs:
+                logger.info(f"sglang_overrides: overriding {key}={kwargs[key]} -> {value} (rank={rank})")
+            kwargs[key] = value
+            unused_keys.discard(key)
 
     # for compatibility with old args
     if len(unused_keys) > 0:
@@ -608,6 +582,5 @@ def _compute_server_args(
     "dist_init_addr",
     "skip_server_warmup",
     "enable_draft_weights_cpu_backup",
-    "enable_metrics",
     "mem_fraction_static",
 ]
diff --git a/slime/backends/vllm_utils/__init__.py b/slime/backends/vllm_utils/__init__.py
deleted file mode 100644
index 41b5f9909c..0000000000
--- a/slime/backends/vllm_utils/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from slime.backends.vllm_utils.vllm_engine import VLLMEngine
-from slime.backends.vllm_utils.vllm_translation_sidecar import TranslationSidecar, run_sidecar
-
-__all__ = ["VLLMEngine", "TranslationSidecar", "run_sidecar"]
diff --git a/slime/backends/vllm_utils/vllm_engine.py b/slime/backends/vllm_utils/vllm_engine.py
deleted file mode 100644
index c2902ea13f..0000000000
--- a/slime/backends/vllm_utils/vllm_engine.py
+++ /dev/null
@@ -1,375 +0,0 @@
-"""VLLMEngine: Ray actor that launches and manages a vLLM server + translation sidecar."""
-
-import logging
-import multiprocessing
-import os
-import subprocess
-import tempfile
-import time
-
-import requests
-
-from slime.ray.ray_actor import RayActor
-from slime.utils.http_utils import get_host_info
-from slime.utils.misc import get_free_port
-
-logger = logging.getLogger(__name__)
-
-
-class VLLMEngine(RayActor):
-    """Ray actor that runs vLLM server with same interface as SGLangEngine for weight sync."""
-
-    def __init__(self, args, rank: int, base_gpu_id: int | None = None, gpu_ids: list[int] | None = None, **kwargs):
-        self.args = args
-        self.rank = rank
-        self.base_gpu_id = base_gpu_id or 0
-        self.gpu_ids = gpu_ids or [self.base_gpu_id]
-        self.server_host = None
-        self.server_port = None
-        self.sidecar_port = None
-        self.process = None
-        self.sidecar_process = None
-        self._log_file = None
-        self._sidecar_log_file = None
-        self._weight_version: int = 0
-        self._colocate: bool = getattr(args, "colocate", False)
-
-    @property
-    def sidecar_url(self) -> str:
-        """URL of the translation sidecar (registered with the router)."""
-        return f"http://{self.server_host}:{self.sidecar_port}"
-
-    @property
-    def vllm_url(self) -> str:
-        """URL of the raw vLLM server."""
-        return f"http://{self.server_host}:{self.server_port}"
-
-    def init(self, port=None, host=None, router_ip=None, router_port=None, **kwargs):
-        self.server_host = host or get_host_info()[1]
-        self.server_port = port or get_free_port(15000)
-        self.sidecar_port = get_free_port(self.server_port + 100)
-        self.router_ip = router_ip or getattr(self.args, "sglang_router_ip", None)
-        self.router_port = router_port or getattr(self.args, "sglang_router_port", None)
-
-        model = getattr(self.args, "vllm_model", None) or self.args.hf_checkpoint
-        self._model_name = model
-        tp = self.args.rollout_num_gpus_per_engine
-        gpu_ids = self.gpu_ids[:tp]
-        cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
-        if cvd:
-            visible = [int(x) for x in cvd.split(",") if x.strip()]
-            dev_str = ",".join(str(visible[gid]) if gid < len(visible) else str(gid) for gid in gpu_ids)
-        else:
-            dev_str = ",".join(str(g) for g in gpu_ids)
-
-        seed = getattr(self.args, "seed", 1234) + self.rank
-        # Use IPC backend when colocated (same GPU), NCCL otherwise
-        if self._colocate:
-            weight_transfer_backend = '{"backend": "ipc"}'
-        else:
-            weight_transfer_backend = '{"backend": "nccl"}'
-
-        cmd = [
-            "vllm", "serve", model,
-            "--tensor-parallel-size", str(tp),
-            "--port", str(self.server_port),
-            "--host", "0.0.0.0",
-            "--weight-transfer-config", weight_transfer_backend,
-            "--seed", str(seed),
-            "--trust-remote-code",
-        ]
-        gpu_mem_util = getattr(self.args, "vllm_gpu_memory_utilization", 0.25)
-        cmd.extend(["--gpu-memory-utilization", str(gpu_mem_util)])
-        if getattr(self.args, "offload_rollout", False):
-            cmd.append("--enable-sleep-mode")
-        if getattr(self.args, "vllm_enforce_eager", True):
-            cmd.append("--enforce-eager")
-        if getattr(self.args, "fp16", False):
-            cmd.extend(["--dtype", "float16"])
-
-        env = os.environ.copy()
-        env["VLLM_SERVER_DEV_MODE"] = "1"
-        if self._colocate:
-            env["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
-        env["CUDA_VISIBLE_DEVICES"] = dev_str
-        env.setdefault("NCCL_DEBUG", "INFO")
-        env.setdefault("NCCL_DEBUG_SUBSYS", "ALL")
-        env["NCCL_P2P_DISABLE"] = "1"
-        env.setdefault("NCCL_IB_DISABLE", "1")
-
-        self._log_file = tempfile.NamedTemporaryFile(
-            prefix="vllm_engine_", suffix=".log", delete=False, mode="w"
-        )
-        
-        logger.info("Launching vLLM: cmd=%s, CUDA_VISIBLE_DEVICES=%s, log=%s",
-                     " ".join(cmd), dev_str, self._log_file.name)
-        self.process = subprocess.Popen(
-            cmd,
-            env=env,
-            stdout=self._log_file,
-            stderr=subprocess.STDOUT,
-        )
-        self._wait_healthy()
-
-        # Launch the translation sidecar
-        self._launch_sidecar()
-        self._wait_sidecar_healthy()
-
-        # Register the sidecar URL with the router
-        if self.router_ip and self.router_port:
-            self._register_with_router()
-
-    def _wait_healthy(self, timeout=300):
-        base = f"http://{self.server_host}:{self.server_port}"
-        start = time.time()
-        while time.time() - start < timeout:
-            try:
-                r = requests.get(f"{base}/health", timeout=5)
-                if r.status_code == 200:
-                    logger.info("vLLM server healthy at %s:%s", self.server_host, self.server_port)
-                    return
-            except Exception:
-                pass
-            if self.process and self.process.poll() is not None:
-                log_tail = self._read_log_tail()
-                raise RuntimeError(f"vLLM process exited with code {self.process.returncode}.\n{log_tail}")
-            time.sleep(2)
-        log_tail = self._read_log_tail()
-        raise TimeoutError(f"vLLM server failed to become healthy within {timeout}s.\n{log_tail}")
-
-    def _launch_sidecar(self):
-        """Launch the translation sidecar as a subprocess."""
-        from slime.backends.vllm_utils.vllm_translation_sidecar import run_sidecar
-
-        self._sidecar_log_file = tempfile.NamedTemporaryFile(
-            prefix="vllm_sidecar_", suffix=".log", delete=False, mode="w"
-        )
-
-        def _target():
-            run_sidecar(
-                vllm_host="127.0.0.1",
-                vllm_port=self.server_port,
-                sidecar_host="0.0.0.0",
-                sidecar_port=self.sidecar_port,
-                model_name=self._model_name,
-                log_level="info",
-            )
-
-        self.sidecar_process = multiprocessing.Process(target=_target, daemon=True)
-        self.sidecar_process.start()
-        logger.info(
-            "Launched translation sidecar on port %s (vLLM → %s:%s), log=%s",
-            self.sidecar_port,
-            self.server_host,
-            self.server_port,
-            self._sidecar_log_file.name,
-        )
-
-    def _wait_sidecar_healthy(self, timeout: float = 60.0):
-        """Block until the sidecar /health endpoint responds 200."""
-        url = f"{self.sidecar_url}/health"
-        start = time.time()
-        while time.time() - start < timeout:
-            try:
-                r = requests.get(url, timeout=5)
-                if r.status_code == 200:
-                    logger.info("Translation sidecar healthy at %s", self.sidecar_url)
-                    return
-            except Exception:
-                pass
-            if self.sidecar_process and not self.sidecar_process.is_alive():
-                raise RuntimeError(
-                    f"Sidecar process exited with code {self.sidecar_process.exitcode}"
-                )
-            time.sleep(1)
-        raise TimeoutError(f"Sidecar failed to become healthy within {timeout}s")
-
-    def _register_with_router(self):
-        """Register the sidecar URL with the SlimeRouter."""
-        router_url = f"http://{self.router_ip}:{self.router_port}"
-        response = requests.post(
-            f"{router_url}/add_worker",
-            params={"url": self.sidecar_url},
-        )
-        response.raise_for_status()
-        logger.info(
-            "Registered sidecar %s with router at %s",
-            self.sidecar_url,
-            router_url,
-        )
-
-    def _bump_weight_version(self, version: int | None = None):
-        """Notify the sidecar to increment (or set) its weight version counter."""
-        url = f"{self.sidecar_url}/set_weight_version"
-        payload = {"weight_version": version} if version is not None else {}
-        try:
-            r = requests.post(url, json=payload, timeout=10)
-            r.raise_for_status()
-            self._weight_version = r.json().get("weight_version", self._weight_version)
-        except Exception as exc:
-            logger.warning("Failed to bump sidecar weight version: %s", exc)
-
-    def _read_log_tail(self, n=200):
-        if not self._log_file:
-            return ""
-        try:
-            self._log_file.flush()
-            with open(self._log_file.name) as f:
-                lines = f.readlines()
-            return "".join(lines[-n:])
-        except Exception:
-            return ""
-
-    def _post(self, path: str, json_data: dict | None = None, params: dict | None = None):
-        url = f"http://{self.server_host}:{self.server_port}{path}"
-        kwargs = {"timeout": 120}
-        if json_data is not None:
-            kwargs["json"] = json_data
-        if params is not None:
-            kwargs["params"] = params
-        r = requests.post(url, **kwargs)
-        if not r.ok:
-            body = r.text[:2000] if r.text else "(empty)"
-            log_tail = self._read_log_tail(50)
-            logger.error(
-                "vLLM %s returned %s: %s\n--- vLLM log tail ---\n%s",
-                path, r.status_code, body, log_tail,
-            )
-        r.raise_for_status()
-        try:
-            return r.json()
-        except Exception:
-            return None
-
-    def health_generate(self, timeout: float = 5.0) -> bool:
-        try:
-            r = requests.get(
-                f"http://{self.server_host}:{self.server_port}/health",
-                timeout=timeout,
-            )
-            r.raise_for_status()
-            return True
-        except requests.RequestException:
-            return False
-
-    def pause_generation(self):
-        self._post("/pause", params={"mode": "abort"})
-
-    def flush_cache(self):
-        pass
-
-    def init_weights_update_group(self, master_address, master_port, rank_offset, world_size, group_name=None, backend=None):
-        logger.info(
-            "Initializing NCCL weight transfer: master=%s:%s, rank_offset=%d, "
-            "world_size=%d, vllm_url=http://%s:%s, vllm_log=%s",
-            master_address, master_port, rank_offset, world_size,
-            self.server_host, self.server_port,
-            self._log_file.name if self._log_file else "<none>",
-        )
-        self._post("/init_weight_transfer_engine", json_data={
-            "init_info": {
-                "master_address": master_address,
-                "master_port": master_port,
-                "rank_offset": rank_offset,
-                "world_size": world_size,
-            }
-        })
-        log_tail = self._read_log_tail(30)
-        logger.info("vLLM log after init_weight_transfer_engine:\n%s", log_tail)
-
-    def update_weights_from_distributed(
-        self,
-        names,
-        dtypes,
-        shapes,
-        group_name=None,
-        flush_cache=False,
-        weight_version=None,
-        packed: bool = True,
-    ):
-        dtype_names = [str(d).replace("torch.", "") for d in dtypes]
-        shape_lists = [list(s) for s in shapes]
-        self._post("/update_weights", json_data={
-            "update_info": {
-                "names": names,
-                "dtype_names": dtype_names,
-                "shapes": shape_lists,
-                "packed": packed,
-            }
-        })
-        # Notify the sidecar about the new weight version
-        self._bump_weight_version(weight_version)
-
-    def continue_generation(self):
-        self._post("/resume")
-
-    def destroy_weights_update_group(self, group_name):
-        pass
-
-    def release_memory_occupation(self):
-        try:
-            self._post("/sleep", params={"level": "1", "mode": "abort"})
-        except requests.RequestException as e:
-            logger.warning("vLLM sleep failed (need --enable-sleep-mode?): %s", e)
-
-    def resume_memory_occupation(self, tags: list[str] | None = None):
-        try:
-            params = {}
-            if tags:
-                params["tags"] = tags
-            self._post("/wake_up", params=params)
-        except requests.RequestException as e:
-            logger.warning("vLLM wake_up failed: %s", e)
-
-    def get_weight_version(self):
-        if self.sidecar_port:
-            try:
-                r = requests.get(f"{self.sidecar_url}/get_weight_version", timeout=5)
-                r.raise_for_status()
-                return r.json().get("weight_version", self._weight_version)
-            except Exception:
-                pass
-        return self._weight_version
-
-    def get_vllm_url(self) -> str:
-        """Return the raw vLLM server URL (used by IPC weight transfer)."""
-        return self.vllm_url
-
-    def set_weight_version(self, version: int) -> None:
-        """Set weight version on both the engine and the sidecar."""
-        self._weight_version = version
-        self._bump_weight_version(version)
-
-    def check_weights(self, action: str):
-        pass
-
-    def post_process_weights(self, **kwargs):
-        pass
-
-    def shutdown(self):
-        # Shutdown translation sidecar first
-        if self.sidecar_process and self.sidecar_process.is_alive():
-            self.sidecar_process.terminate()
-            self.sidecar_process.join(timeout=10)
-            if self.sidecar_process.is_alive():
-                self.sidecar_process.kill()
-            self.sidecar_process = None
-        if self._sidecar_log_file:
-            try:
-                self._sidecar_log_file.close()
-            except Exception:
-                pass
-
-        # Shutdown vLLM server
-        if self.process:
-            self.process.terminate()
-            try:
-                self.process.wait(timeout=10)
-            except subprocess.TimeoutExpired:
-                self.process.kill()
-            self.process = None
-        if self._log_file:
-            try:
-                self._log_file.close()
-            except Exception:
-                pass
diff --git a/slime/backends/vllm_utils/vllm_translation_sidecar.py b/slime/backends/vllm_utils/vllm_translation_sidecar.py
deleted file mode 100644
index 73998d8bf4..0000000000
--- a/slime/backends/vllm_utils/vllm_translation_sidecar.py
+++ /dev/null
@@ -1,454 +0,0 @@
-"""
-vLLM Translation Sidecar
-========================
-
-Lightweight FastAPI process co-located with each vLLM engine.
-Receives SGLang-format ``/generate`` requests from the SlimeRouter,
-translates them to vLLM ``/v1/completions``, and translates responses back.
-
-Also proxies lifecycle endpoints:
-  /health, /health_generate, /abort_request, /flush_cache, /get_weight_version
-
-See docs/en/vllm/ROUTER_DESIGN.md for the full specification.
-"""
-
-from __future__ import annotations
-
-import argparse
-import asyncio
-import logging
-import signal
-import sys
-from contextlib import asynccontextmanager
-from typing import Any
-
-import httpx
-import uvicorn
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-
-logger = logging.getLogger(__name__)
-
-# ---------------------------------------------------------------------------
-# Sampling-param translation tables
-# ---------------------------------------------------------------------------
-
-# SGLang name → vLLM /v1/completions name  (only entries that differ)
-_PARAM_RENAME = {
-    "max_new_tokens": "max_tokens",
-    "no_stop_trim": "include_stop_str_in_output",
-    "sampling_seed": "seed",
-}
-
-# Parameters passed through unchanged
-_PARAM_DIRECT = frozenset(
-    {
-        "temperature",
-        "top_p",
-        "top_k",
-        "stop",
-        "stop_token_ids",
-        "skip_special_tokens",
-        "spaces_between_special_tokens",
-    }
-)
-
-# finish_reason vLLM → SGLang-style {"type": ...}
-_FINISH_REASON_MAP = {
-    "stop": "stop",
-    "length": "length",
-    None: "abort",
-}
-
-
-# ---------------------------------------------------------------------------
-# Request / response translation helpers
-# ---------------------------------------------------------------------------
-
-
-def translate_generate_request(
-    body: dict[str, Any],
-    model_name: str,
-) -> dict[str, Any]:
-    """Translate an SGLang-format /generate request → vLLM /v1/completions payload."""
-
-    sp: dict = body.get("sampling_params", {})
-
-    vllm_payload: dict[str, Any] = {
-        "model": model_name,
-        # vLLM accepts list[int] as a pre-tokenized prompt
-        "prompt": body.get("input_ids") or body.get("input_tokens", []),
-        "stream": False,
-    }
-
-    # --- direct-copy params ---
-    for key in _PARAM_DIRECT:
-        if key in sp:
-            vllm_payload[key] = sp[key]
-
-    # --- renamed params ---
-    for src, dst in _PARAM_RENAME.items():
-        if src in sp:
-            vllm_payload[dst] = sp[src]
-
-    # --- logprob handling ---
-    if body.get("return_logprob", False):
-        vllm_payload["logprobs"] = 1
-        # request token IDs alongside logprobs
-        # NOTE: must be a top-level param; "extra_body" is an OpenAI SDK
-        # client concept and is ignored by vLLM's raw HTTP API.
-        vllm_payload["return_token_ids"] = True
-
-    return vllm_payload
-
-
-def translate_vllm_response(
-    vllm_resp: dict[str, Any],
-    weight_version: int,
-) -> dict[str, Any]:
-    """Translate a vLLM /v1/completions response → SGLang-format JSON."""
-
-    choice: dict = vllm_resp.get("choices", [{}])[0]
-    usage: dict = vllm_resp.get("usage", {})
-
-    # --- output token IDs ---
-    output_ids: list[int] = choice.get("token_ids", [])
-
-    # --- logprobs: zip(logprob, token_id) ---
-    output_token_logprobs: list[list[float | int]] = []
-    logprobs_obj = choice.get("logprobs")
-    if logprobs_obj and output_ids:
-        raw_lp: list[float | None] = logprobs_obj.get("token_logprobs", [])
-        output_token_logprobs = [
-            [float(lp) if lp is not None else 0.0, tid]
-            for lp, tid in zip(raw_lp, output_ids)
-        ]
-
-    # --- finish reason ---
-    raw_reason = choice.get("finish_reason")
-    mapped = _FINISH_REASON_MAP.get(raw_reason, raw_reason or "abort")
-    finish_reason = {"type": mapped}
-
-    meta_info: dict[str, Any] = {
-        "finish_reason": finish_reason,
-        "weight_version": weight_version,
-        "prompt_tokens": usage.get("prompt_tokens", 0),
-        "completion_tokens": usage.get("completion_tokens", len(output_ids)),
-        "cached_tokens": 0,
-    }
-    # Only include output_token_logprobs when we have valid paired data;
-    # a None value causes RadixTreeMiddleware to silently fail when iterating.
-    if output_token_logprobs:
-        meta_info["output_token_logprobs"] = output_token_logprobs
-
-    return {
-        "text": choice.get("text", ""),
-        "output_ids": output_ids,
-        "meta_info": meta_info,
-    }
-
-
-# ---------------------------------------------------------------------------
-# Sidecar application
-# ---------------------------------------------------------------------------
-
-
-class TranslationSidecar:
-    """Manages state and provides the FastAPI app for the translation sidecar."""
-
-    def __init__(
-        self,
-        vllm_base_url: str,
-        model_name: str,
-        *,
-        timeout: float = 600.0,
-        max_connections: int = 256,
-    ):
-        self.vllm_base_url = vllm_base_url.rstrip("/")
-        self.model_name = model_name
-        self._weight_version: int = 0
-        self._active_connections: set[httpx.Response] = set()
-        self._lock = asyncio.Lock()
-
-        self._client: httpx.AsyncClient | None = None
-        self._timeout = timeout
-        self._max_connections = max_connections
-
-        self.app = self._build_app()
-
-    # ---- lifecycle -----------------------------------------------------------
-
-    async def startup(self):
-        self._client = httpx.AsyncClient(
-            limits=httpx.Limits(
-                max_connections=self._max_connections,
-                max_keepalive_connections=self._max_connections,
-            ),
-            timeout=httpx.Timeout(self._timeout),
-        )
-
-    async def shutdown(self):
-        if self._client:
-            await self._client.aclose()
-            self._client = None
-
-    # ---- app factory ---------------------------------------------------------
-
-    def _build_app(self) -> FastAPI:
-
-        @asynccontextmanager
-        async def lifespan(app: FastAPI):
-            await self.startup()
-            yield
-            await self.shutdown()
-
-        app = FastAPI(title="vLLM Translation Sidecar", lifespan=lifespan)
-
-        app.post("/generate")(self.generate)
-        app.get("/health")(self.health)
-        app.get("/health_generate")(self.health_generate)
-        app.post("/abort_request")(self.abort_request)
-        app.get("/flush_cache")(self.flush_cache)
-        app.get("/get_weight_version")(self.get_weight_version)
-        app.post("/set_weight_version")(self.set_weight_version)
-
-        return app
-
-    # ---- endpoints -----------------------------------------------------------
-
-    async def generate(self, request: Request):
-        """Translate SGLang /generate → vLLM /v1/completions → SGLang response."""
-
-        body = await request.json()
-        vllm_payload = translate_generate_request(body, self.model_name)
-
-        url = f"{self.vllm_base_url}/v1/completions"
-
-        resp: httpx.Response | None = None
-        try:
-            async with self._lock:
-                # We don't actually hold the lock during the request,
-                # just use it to safely add to the tracking set.
-                pass
-
-            resp = await self._client.post(url, json=vllm_payload)
-            self._active_connections.add(resp)
-
-            resp.raise_for_status()
-            vllm_data = resp.json()
-        except httpx.HTTPStatusError as exc:
-            logger.error(
-                "vLLM /v1/completions returned %s: %s",
-                exc.response.status_code,
-                exc.response.text[:2000],
-            )
-            # Return an abort-style response so the middleware retries
-            return JSONResponse(
-                content={
-                    "text": "",
-                    "output_ids": [],
-                    "meta_info": {
-                        "finish_reason": {"type": "abort"},
-                        "weight_version": self._weight_version,
-                        "prompt_tokens": 0,
-                        "completion_tokens": 0,
-                        "cached_tokens": 0,
-                    },
-                },
-                status_code=200,
-            )
-        except Exception as exc:
-            logger.error("Error calling vLLM: %s", exc, exc_info=True)
-            return JSONResponse(
-                content={
-                    "text": "",
-                    "output_ids": [],
-                    "meta_info": {
-                        "finish_reason": {"type": "abort"},
-                        "weight_version": self._weight_version,
-                        "prompt_tokens": 0,
-                        "completion_tokens": 0,
-                        "cached_tokens": 0,
-                    },
-                },
-                status_code=200,
-            )
-        finally:
-            if resp is not None:
-                self._active_connections.discard(resp)
-                await resp.aclose()
-
-        translated = translate_vllm_response(vllm_data, self._weight_version)
-        return JSONResponse(content=translated)
-
-    async def health(self):
-        """Proxy to vLLM's built-in /health endpoint."""
-        try:
-            resp = await self._client.get(f"{self.vllm_base_url}/health", timeout=5.0)
-            return JSONResponse(content={"status": "ok"}, status_code=resp.status_code)
-        except Exception:
-            return JSONResponse(content={"status": "unhealthy"}, status_code=503)
-
-    async def health_generate(self):
-        """
-        Startup readiness probe.
-
-        Checks vLLM /health and optionally fires a dummy /v1/completions
-        with max_tokens=1 to verify end-to-end readiness.
-        """
-        try:
-            resp = await self._client.get(f"{self.vllm_base_url}/health", timeout=10.0)
-            if resp.status_code != 200:
-                return JSONResponse(content={"status": "not_ready"}, status_code=503)
-
-            # Lightweight smoke test: single-token completion
-            dummy_payload = {
-                "model": self.model_name,
-                "prompt": "hi",
-                "max_tokens": 1,
-                "stream": False,
-            }
-            resp2 = await self._client.post(
-                f"{self.vllm_base_url}/v1/completions",
-                json=dummy_payload,
-                timeout=30.0,
-            )
-            if resp2.status_code == 200:
-                return JSONResponse(content={"status": "ready"}, status_code=200)
-            else:
-                return JSONResponse(content={"status": "not_ready"}, status_code=503)
-        except Exception as exc:
-            logger.debug("health_generate check failed: %s", exc)
-            return JSONResponse(content={"status": "not_ready"}, status_code=503)
-
-    async def abort_request(self, request: Request):
-        """
-        Handle abort requests.
-
-        vLLM uses HTTP disconnect for cancellation.  We close all active
-        connections to the vLLM backend, which triggers vLLM's
-        ``@with_cancellation`` decorator to abort in-flight requests.
-
-        Alternatively, use pause/resume for a cleaner between-generation abort.
-        """
-        body = await request.json()
-
-        if body.get("abort_all", False):
-            # Strategy: pause + resume clears the pipeline
-            try:
-                await self._client.post(
-                    f"{self.vllm_base_url}/pause",
-                    params={"mode": "abort"},
-                    timeout=30.0,
-                )
-                await self._client.post(
-                    f"{self.vllm_base_url}/resume",
-                    timeout=30.0,
-                )
-            except Exception as exc:
-                logger.warning("pause/resume abort failed, falling back to connection close: %s", exc)
-                # Fallback: close all tracked connections
-                conns = list(self._active_connections)
-                self._active_connections.clear()
-                for conn in conns:
-                    try:
-                        await conn.aclose()
-                    except Exception:
-                        pass
-
-        return JSONResponse(content={"status": "ok"})
-
-    async def flush_cache(self):
-        """
-        Flush the KV cache.
-
-        vLLM equivalent: sleep(level=1) + wake_up(tags=kv_cache).
-        """
-        try:
-            await self._client.post(
-                f"{self.vllm_base_url}/sleep",
-                params={"level": "1", "mode": "abort"},
-                timeout=30.0,
-            )
-            await self._client.post(
-                f"{self.vllm_base_url}/wake_up",
-                params={"tags": "kv_cache"},
-                timeout=30.0,
-            )
-            return JSONResponse(content={"status": "ok"})
-        except Exception as exc:
-            logger.warning("flush_cache failed: %s", exc)
-            return JSONResponse(content={"status": "ok"})
-
-    async def get_weight_version(self):
-        """Return the sidecar-tracked weight version counter."""
-        return JSONResponse(content={"weight_version": self._weight_version})
-
-    async def set_weight_version(self, request: Request):
-        """Increment or set the weight version (called by VLLMEngine after weight update)."""
-        body = await request.json()
-        if "weight_version" in body:
-            self._weight_version = int(body["weight_version"])
-        else:
-            self._weight_version += 1
-        return JSONResponse(content={"weight_version": self._weight_version})
-
-
-# ---------------------------------------------------------------------------
-# Standalone entry point
-# ---------------------------------------------------------------------------
-
-
-def run_sidecar(
-    vllm_host: str = "127.0.0.1",
-    vllm_port: int = 8000,
-    sidecar_host: str = "0.0.0.0",
-    sidecar_port: int = 8100,
-    model_name: str = "default",
-    timeout: float = 600.0,
-    max_connections: int = 256,
-    log_level: str = "info",
-):
-    """Launch the translation sidecar as a standalone uvicorn process."""
-
-    vllm_base_url = f"http://{vllm_host}:{vllm_port}"
-    sidecar = TranslationSidecar(
-        vllm_base_url=vllm_base_url,
-        model_name=model_name,
-        timeout=timeout,
-        max_connections=max_connections,
-    )
-    uvicorn.run(
-        sidecar.app,
-        host=sidecar_host,
-        port=sidecar_port,
-        log_level=log_level,
-    )
-
-
-def main():
-    parser = argparse.ArgumentParser(description="vLLM Translation Sidecar")
-    parser.add_argument("--vllm-host", type=str, default="127.0.0.1")
-    parser.add_argument("--vllm-port", type=int, default=8000)
-    parser.add_argument("--sidecar-host", type=str, default="0.0.0.0")
-    parser.add_argument("--sidecar-port", type=int, default=8100)
-    parser.add_argument("--model-name", type=str, default="default")
-    parser.add_argument("--timeout", type=float, default=600.0)
-    parser.add_argument("--max-connections", type=int, default=256)
-    parser.add_argument("--log-level", type=str, default="info")
-    args = parser.parse_args()
-
-    run_sidecar(
-        vllm_host=args.vllm_host,
-        vllm_port=args.vllm_port,
-        sidecar_host=args.sidecar_host,
-        sidecar_port=args.sidecar_port,
-        model_name=args.model_name,
-        timeout=args.timeout,
-        max_connections=args.max_connections,
-        log_level=args.log_level,
-    )
-
-
-if __name__ == "__main__":
-    main()
diff --git a/slime/ray/actor_group.py b/slime/ray/actor_group.py
index d50a61d1c4..5542835313 100644
--- a/slime/ray/actor_group.py
+++ b/slime/ray/actor_group.py
@@ -50,45 +50,16 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
         assert pg is not None
         pg, reordered_bundle_indices, _reordered_gpu_ids = pg
 
-        # Restrict CUDA_VISIBLE_DEVICES to only this group's GPUs so that
-        # NCCL / PyTorch do not allocate memory on rollout GPUs.
-        trainer_gpu_ids = [_reordered_gpu_ids[rank] for rank in range(world_size)]
-        trainer_cvd = ",".join(str(g) for g in trainer_gpu_ids)
-
         env_vars = {
             # because sglang will always set NCCL_CUMEM_ENABLE to 0
             # we need also set it to 0 to prevent nccl error.
             "NCCL_CUMEM_ENABLE": os.environ.get("NCCL_CUMEM_ENABLE", "0"),
             "NVTE_FP8_BLOCK_SCALING_FP32_SCALES": os.environ.get("NVTE_FP8_BLOCK_SCALING_FP32_SCALES", "1"),
             **{name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST},
-            "CUDA_VISIBLE_DEVICES": trainer_cvd,
             **self.args.train_env_vars,
         }
 
-        # if self.args.offload_train and self.args.train_backend == "megatron":
-        #     import torch_memory_saver
-
-        #     dynlib_path = os.path.join(
-        #         os.path.dirname(os.path.dirname(torch_memory_saver.__file__)),
-        #         "torch_memory_saver_hook_mode_preload.abi3.so",
-        #     )
-        #     assert os.path.exists(dynlib_path), f"LD_PRELOAD so file {dynlib_path} does not exist."
-
-        #     env_vars["LD_PRELOAD"] = dynlib_path
-        #     env_vars["TMS_INIT_ENABLE"] = "1"
-        #     env_vars["TMS_INIT_ENABLE_CPU_BACKUP"] = "1"
-
-        # torch_memory_saver requires an LD_PRELOAD'd allocator hook. It is only
-        # used when offloading the trainer for the sglang rollout backend. With
-        # the vLLM rollout backend we instead fall back to an explicit CPU
-        # offload path inside the Megatron actor (see ``MegatronTrainRayActor``
-        # ``sleep`` / ``wake_up``), so we must NOT preload the .so here.
-        rollout_backend = getattr(self.args, "rollout_backend", "sglang")
-        if (
-            self.args.offload_train
-            and self.args.train_backend == "megatron"
-            and rollout_backend != "vllm"
-        ):
+        if self.args.offload_train and self.args.train_backend == "megatron":
             import torch_memory_saver
 
             dynlib_path = os.path.join(
@@ -105,9 +76,16 @@ def _allocate_gpus_for_actor(self, pg, num_gpus_per_actor):
         if self.args.use_routing_replay and self.role == "actor":
             env_vars["ENABLE_ROUTING_REPLAY"] = "1"
 
-        from slime.backends.megatron_utils.actor import MegatronTrainRayActor
+        backend = self.args.train_backend
+        if backend == "megatron":
+            from slime.backends.megatron_utils.actor import MegatronTrainRayActor
+
+            actor_impl = MegatronTrainRayActor
+
+        else:
+            from slime.backends.fsdp_utils import FSDPTrainRayActor
 
-        actor_impl = MegatronTrainRayActor
+            actor_impl = FSDPTrainRayActor
 
         TrainRayActor = ray.remote(num_gpus=1, runtime_env={"env_vars": env_vars})(actor_impl)
 
@@ -137,25 +115,9 @@ def async_init(self, args, role, with_ref=False, with_opd_teacher=False):
             for actor in self._actor_handlers
         ]
 
-    def async_train(self, rollout_id, rollout_data_ref, external_data=None):
-        """Do one rollout training. Returns a list of Ray refs (one per worker).
-
-        For critics, each ref resolves to ``{"values": [cpu tensors...]}`` (or ``{}``
-        for non-last-PP-stage workers). Actor refs resolve to ``None``.
-
-        ``external_data`` may be a list (one item per worker) or a single dict
-        broadcast to all workers.
-        """
-        if isinstance(external_data, list):
-            assert len(external_data) == len(self._actor_handlers)
-            return [
-                actor.train.remote(rollout_id, rollout_data_ref, external_data=ed)
-                for actor, ed in zip(self._actor_handlers, external_data, strict=False)
-            ]
-        return [
-            actor.train.remote(rollout_id, rollout_data_ref, external_data=external_data)
-            for actor in self._actor_handlers
-        ]
+    def async_train(self, rollout_id, rollout_data_ref):
+        """Do one rollout training"""
+        return [actor.train.remote(rollout_id, rollout_data_ref) for actor in self._actor_handlers]
 
     def save_model(self, rollout_id, force_sync=False):
         """Save actor model"""
@@ -174,5 +136,13 @@ def offload(self):
     def clear_memory(self):
         return ray.get([actor.clear_memory.remote() for actor in self._actor_handlers])
 
+    def connect(self, critic_group):
+        return ray.get(
+            [
+                actor.connect_actor_critic.remote(critic)
+                for actor, critic in zip(self._actor_handlers, critic_group._actor_handlers, strict=False)
+            ]
+        )
+
     def set_rollout_manager(self, rollout_manager):
         return ray.get([actor.set_rollout_manager.remote(rollout_manager) for actor in self._actor_handlers])
diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py
index edc723dcef..ef78e271eb 100644
--- a/slime/ray/placement_group.py
+++ b/slime/ray/placement_group.py
@@ -77,93 +77,91 @@ def _create_placement_group(num_gpus):
 
 
 def create_placement_groups(args):
-    """Create placement groups for actor, critic, and rollout engines."""
+    """Create placement groups for actor and rollout engines."""
 
     num_gpus = 0
     if args.debug_train_only:
         num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node
         rollout_offset = 0
+        if args.use_critic:
+            num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node
+            critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node
     elif args.debug_rollout_only:
         num_gpus = args.rollout_num_gpus
         rollout_offset = 0
     elif args.colocate:
         num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node
         rollout_offset = 0
+        if args.use_critic:
+            num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node
+            critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node
     else:
         num_gpus = args.actor_num_nodes * args.actor_num_gpus_per_node + args.rollout_num_gpus
         rollout_offset = args.actor_num_nodes * args.actor_num_gpus_per_node
+        if args.use_critic:
+            num_gpus += args.critic_num_nodes * args.critic_num_gpus_per_node
+            critic_offset = args.actor_num_nodes * args.actor_num_gpus_per_node
+            rollout_offset += args.critic_num_nodes * args.critic_num_gpus_per_node
 
     logger.info(f"Creating placement group with {num_gpus} GPUs...")
     pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids = _create_placement_group(num_gpus)
+
     rollout_pg_reordered_bundle_indices = actor_pg_reordered_bundle_indices[rollout_offset:]
     rollout_pg_reordered_gpu_ids = actor_pg_reordered_gpu_ids[rollout_offset:]
+    if args.use_critic:
+        critic_pg_reordered_bundle_indices = actor_pg_reordered_bundle_indices[critic_offset:]
+        critic_pg_reordered_gpu_ids = actor_pg_reordered_gpu_ids[critic_offset:]
 
-    result = {
+    return {
         "actor": (pg, actor_pg_reordered_bundle_indices, actor_pg_reordered_gpu_ids),
+        "critic": (pg, critic_pg_reordered_bundle_indices, critic_pg_reordered_gpu_ids) if args.use_critic else None,
         "rollout": (pg, rollout_pg_reordered_bundle_indices, rollout_pg_reordered_gpu_ids),
     }
 
-    result["critic"] = result["actor"] if args.use_critic else None
-
-    return result
-
 
-def allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="actor"):
+def allocate_train_group(args, num_nodes, num_gpus_per_node, pg):
     return RayTrainGroup(
         args=args,
         num_nodes=num_nodes,
         num_gpus_per_node=num_gpus_per_node,
         pg=pg,
         num_gpus_per_actor=0.4,
-        role=role,
     )
 
 
 def create_training_models(args, pgs, rollout_manager):
-    actor_args = args
-    if args.megatron_config_path is not None:
-        from slime.utils.arguments import parse_megatron_role_args
-
-        actor_args = parse_megatron_role_args(args, args.megatron_config_path, role="actor")
-
     actor_model = allocate_train_group(
-        args=actor_args,
+        args=args,
         num_nodes=args.actor_num_nodes,
         num_gpus_per_node=args.actor_num_gpus_per_node,
         pg=pgs["actor"],
     )
-
-    critic_model = None
     if args.use_critic:
-        from slime.utils.arguments import parse_megatron_role_args
-
-        critic_args = (
-            parse_megatron_role_args(args, args.megatron_config_path, role="critic")
-            if args.megatron_config_path is not None
-            else args
-        )
         critic_model = allocate_train_group(
-            args=critic_args,
+            args=args,
             num_nodes=args.critic_num_nodes,
             num_gpus_per_node=args.critic_num_gpus_per_node,
             pg=pgs["critic"],
-            role="critic",
         )
-        critic_start_rollout_ids = ray.get(critic_model.async_init(critic_model.args, role="critic", with_ref=False))
+        critic_init_handle = critic_model.async_init(args, role="critic", with_ref=False)
+    else:
+        critic_model = None
 
-    actor_start_rollout_ids = ray.get(
+    start_rollout_ids = ray.get(
         actor_model.async_init(
-            actor_args,
+            args,
             role="actor",
-            with_ref=actor_args.kl_coef != 0 or actor_args.use_kl_loss,
-            with_opd_teacher=actor_args.use_opd and actor_args.opd_type == "megatron",
+            with_ref=args.kl_coef != 0 or args.use_kl_loss,
+            with_opd_teacher=args.use_opd and args.opd_type == "megatron",
         )
     )
-    # TODO how to decide rollout start id when critic is involved? For now we just require user to specify it via args.
+
     if args.use_critic:
-        start_rollout_ids = critic_start_rollout_ids
-    else:
-        start_rollout_ids = actor_start_rollout_ids
+        critic_start_rollout_ids = ray.get(critic_init_handle)
+        if not args.critic_train_only:
+            actor_model.connect(critic_model)
+        else:
+            start_rollout_ids = critic_start_rollout_ids
 
     assert len(set(start_rollout_ids)) == 1
 
diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py
index fbb6338f93..8f919c1172 100644
--- a/slime/ray/rollout.py
+++ b/slime/ray/rollout.py
@@ -11,15 +11,11 @@
 import numpy as np
 import ray
 import torch
+import yaml
 from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH, GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_WEIGHTS
 
-# GPU memory-tag constants (originally from sglang.srt.constants).
-# Duplicated here to avoid a hard sglang dependency when running the vLLM backend.
-GPU_MEMORY_TYPE_KV_CACHE = "kv_cache"
-GPU_MEMORY_TYPE_WEIGHTS = "weights"
-GPU_MEMORY_TYPE_CUDA_GRAPH = "cuda_graph"
-
-from slime.backends.sglang_utils.sglang_config import ModelConfig, ServerGroupConfig, SglangConfig
+from slime.backends.sglang_utils.sglang_engine import SGLangEngine
 from slime.rollout.base_types import call_rollout_fn
 from slime.utils import logging_utils
 from slime.utils.health_monitor import RolloutHealthMonitor
@@ -40,11 +36,158 @@
 
 
 @dataclasses.dataclass
-class ServerGroup:
+class EngineGroupConfig:
+    """Configuration for a single engine group.
+
+    Attributes:
+        worker_type: One of "regular", "prefill", "decode", or "placeholder".
+                     "placeholder" reserves GPU slots without creating engines.
+        num_gpus: Total number of GPUs for this group.
+        num_gpus_per_engine: GPUs per engine for this group.  Overrides the
+                             model-level or global ``--rollout-num-gpus-per-engine``.
+        overrides: Optional dict of SGLang ``ServerArgs`` field overrides.
+                   These are applied on top of the base CLI ``--sglang-*``
+                   arguments in ``_compute_server_args``.
+    """
+
+    worker_type: str
+    num_gpus: int
+    num_gpus_per_engine: int | None = None
+    overrides: dict = dataclasses.field(default_factory=dict)
+
+    def __post_init__(self):
+        valid_types = {"regular", "prefill", "decode", "placeholder"}
+        assert (
+            self.worker_type in valid_types
+        ), f"Invalid worker_type '{self.worker_type}', must be one of {valid_types}"
+        assert self.num_gpus > 0, f"num_gpus must be > 0, got {self.num_gpus}"
+
+
+@dataclasses.dataclass
+class ModelConfig:
+    """Configuration for a single model deployment.
+
+    Attributes:
+        name: Unique name for this model (e.g. "actor", "reward").
+        model_path: HF checkpoint path.  Falls back to ``args.hf_checkpoint``.
+        num_gpus_per_engine: Default GPUs per engine for all groups in this
+                             model.  Individual groups can override.
+        engine_groups: Engine group configurations for this model.
+    """
+
+    name: str
+    model_path: str | None = None
+    num_gpus_per_engine: int | None = None
+    engine_groups: list[EngineGroupConfig] = dataclasses.field(default_factory=list)
+
+    def resolve(self, args) -> None:
+        """Resolve per-group defaults from model-level then args-level values."""
+        default_gpus_per_engine = self.num_gpus_per_engine or args.rollout_num_gpus_per_engine
+        default_model_path = self.model_path or args.hf_checkpoint
+        for g in self.engine_groups:
+            if g.num_gpus_per_engine is None:
+                g.num_gpus_per_engine = default_gpus_per_engine
+            # Inject model_path into overrides so _compute_server_args picks it up.
+            if "model_path" not in g.overrides:
+                g.overrides["model_path"] = default_model_path
+
+    @property
+    def has_pd_disaggregation(self) -> bool:
+        return any(g.worker_type in ("prefill", "decode") for g in self.engine_groups)
+
+    @property
+    def total_num_gpus(self) -> int:
+        return sum(g.num_gpus for g in self.engine_groups)
+
+
+@dataclasses.dataclass
+class SglangConfig:
+    """Configuration for SGLang engine deployment.
+
+    Loaded from ``--sglang-config`` YAML file.
+
+    **Config format**::
+
+        sglang:
+          - name: actor
+            model_path: /path/to/actor
+            num_gpus_per_engine: 2
+            engine_groups:
+              - worker_type: prefill
+                num_gpus: 4
+                num_gpus_per_engine: 2
+              - worker_type: decode
+                num_gpus: 8
+                num_gpus_per_engine: 4
+          - name: reward
+            model_path: /path/to/reward
+            engine_groups:
+              - worker_type: regular
+                num_gpus: 4
+
+    Each model gets its own router.  ``placeholder`` groups reserve GPU
+    slots without creating engines.  ``overrides`` are ``ServerArgs``
+    field names applied on top of the base ``--sglang-*`` CLI args.
+    """
+
+    models: list[ModelConfig]
+
+    @staticmethod
+    def from_yaml(path: str) -> "SglangConfig":
+        with open(path) as f:
+            data = yaml.safe_load(f)
+
+        assert "sglang" in data, (
+            f"sglang config must have a 'sglang' key, got {list(data.keys())}. "
+            f"Wrap your engine_groups inside a model entry under 'sglang'."
+        )
+        models = []
+        for m in data["sglang"]:
+            groups = [EngineGroupConfig(**g) for g in m.get("engine_groups", [])]
+            models.append(
+                ModelConfig(
+                    name=m["name"],
+                    model_path=m.get("model_path"),
+                    num_gpus_per_engine=m.get("num_gpus_per_engine"),
+                    engine_groups=groups,
+                )
+            )
+        return SglangConfig(models=models)
+
+    @staticmethod
+    def from_prefill_num_servers(args) -> "SglangConfig":
+        """Build a config equivalent to the legacy --prefill-num-servers flag."""
+        total_gpus = args.rollout_num_gpus
+        prefill_gpus = args.prefill_num_servers * args.rollout_num_gpus_per_engine
+        decode_gpus = total_gpus - prefill_gpus
+        assert decode_gpus > 0, f"No decode GPUs: total {total_gpus}, prefill {prefill_gpus}"
+        return SglangConfig(
+            models=[
+                ModelConfig(
+                    name="default",
+                    engine_groups=[
+                        EngineGroupConfig(worker_type="prefill", num_gpus=prefill_gpus),
+                        EngineGroupConfig(worker_type="decode", num_gpus=decode_gpus),
+                    ],
+                )
+            ]
+        )
+
+    @property
+    def has_pd_disaggregation(self) -> bool:
+        return any(m.has_pd_disaggregation for m in self.models)
+
+    @property
+    def total_num_gpus(self) -> int:
+        return sum(m.total_num_gpus for m in self.models)
+
+
+@dataclasses.dataclass
+class EngineGroup:
     """A group of homogeneous SGLang engines with the same configuration.
 
     All engines in a group share the same tp_size / nodes_per_engine / pg.
-    A RolloutServer may contain multiple ServerGroups (e.g. prefill vs decode
+    A RolloutServer may contain multiple EngineGroups (e.g. prefill vs decode
     in PD disaggregation).
     """
 
@@ -53,12 +196,10 @@ class ServerGroup:
     all_engines: list
     num_gpus_per_engine: int
     num_new_engines: int
-    worker_type: str = "regular"  # "regular", "prefill", "decode", or "placeholder"
+    worker_type: str = "regular"  # "regular", "prefill", or "decode"
     rank_offset: int = 0  # cumulative engine count before this group
     gpu_offset: int = 0  # cumulative GPU count before this group
     sglang_overrides: dict = dataclasses.field(default_factory=dict)
-    needs_offload: bool = False  # True when this group's GPUs overlap with megatron
-    model_path: str | None = None  # checkpoint path for update_weights_from_disk
     router_ip: str | None = None
     router_port: int | None = None
 
@@ -77,7 +218,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
         Returns ``(init_handles, port_cursors)`` where *init_handles* is a list
         of Ray ObjectRefs and *port_cursors* maps node index → next free port.
         The caller should ``ray.get()`` on the handles to block until the
-        engines are healthy, and pass *port_cursors* to the next server group
+        engines are healthy, and pass *port_cursors* to the next engine group
         so that different groups on the same node don't race for ports.
 
         Placeholder groups (worker_type="placeholder") skip engine creation entirely.
@@ -92,14 +233,7 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
 
         pg, reordered_bundle_indices, reordered_gpu_ids = self.pg
 
-        if getattr(self.args, "rollout_backend", "sglang") == "vllm":
-            from slime.backends.vllm_utils.vllm_engine import VLLMEngine
-
-            RolloutRayActor = ray.remote(VLLMEngine)
-        else:
-            from slime.backends.sglang_utils.sglang_engine import SGLangEngine
-
-            RolloutRayActor = ray.remote(SGLangEngine)
+        RolloutRayActor = ray.remote(SGLangEngine)
 
         rollout_engines = []
         for i in range(len(self.all_engines)):
@@ -123,17 +257,16 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
             env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST} | {
                 key: os.environ.get(key, default_val)
                 for key, default_val in {
-                    "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "true",
-                    "SGLANG_JIT_DEEPGEMM_FAST_WARMUP": "true",
+                    "SGLANG_JIT_DEEPGEMM_PRECOMPILE": "false",
                     "SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true",
                     "SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK": "true",
                     "SGLANG_MEMORY_SAVER_CUDA_GRAPH": "true",
                     "SGLANG_BATCH_INVARIANT_OPS_ENABLE_MM_FALLBACK_VARIANT": "true",
                     "SGLANG_ENABLE_HEALTH_ENDPOINT_GENERATION": "false",
                     "SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE": "false",
-                    "SLIME_ENABLE_PROFILING": "true",
                 }.items()
             }
+
             rollout_engine = RolloutRayActor.options(
                 num_cpus=num_cpus,
                 num_gpus=num_gpus,
@@ -188,74 +321,55 @@ def start_engines(self, port_cursors: dict[int, int] | None = None) -> tuple[lis
     def offload(self):
         """Fire release_memory_occupation on all engines (non-blocking).
 
-        Returns a list of Ray ObjectRefs.  Skipped for groups that do not
-        overlap with megatron GPUs (``needs_offload=False``).
+        Returns a list of Ray ObjectRefs.
         """
-        if not self.needs_offload:
-            return []
         return [engine.release_memory_occupation.remote() for engine in self.engines if engine is not None]
 
     def onload(self, tags: list[str] | None = None):
         """Fire resume_memory_occupation on all engines (non-blocking).
 
-        Returns a list of Ray ObjectRefs.  Skipped for groups that do not
-        overlap with megatron GPUs (``needs_offload=False``).
+        Returns a list of Ray ObjectRefs.
         """
-        if not self.needs_offload:
-            return []
         return [engine.resume_memory_occupation.remote(tags=tags) for engine in self.engines if engine is not None]
 
-    def onload_weights_from_disk(self):
-        """Reload weights from ``model_path`` for non-updatable groups.
-
-        Used instead of ``resume_memory_occupation(tags=[WEIGHTS])`` so that
-        CPU memory is not consumed by offloaded weight copies.
-        """
-        if not self.needs_offload or not self.model_path:
-            return []
-        return [
-            engine.update_weights_from_disk.remote(self.model_path) for engine in self.engines if engine is not None
-        ]
-
 
 @dataclasses.dataclass
 class RolloutServer:
-    """A model served behind a shared router, with one or more server groups.
+    """A model served behind a shared router, with one or more engine groups.
 
     Each RolloutServer represents one model deployed behind a single router.
-    A server may contain multiple ServerGroups with different
+    A server may contain multiple EngineGroups with different
     ``num_gpus_per_engine`` (e.g. prefill TP=2, decode TP=4).
     """
 
-    server_groups: list[ServerGroup]
+    engine_groups: list[EngineGroup]
     router_ip: str | None = None
     router_port: int | None = None
     model_name: str = "default"
-    update_weights: bool = True
 
     @property
     def engines(self):
         """All node-0 engines across all groups (placeholder groups contribute nothing)."""
-        return [e for g in self.server_groups for e in g.engines]
+        return [e for g in self.engine_groups for e in g.engines]
 
     @property
     def all_engines(self):
         """All engines (including non-node-0) across all groups."""
-        return [e for g in self.server_groups for e in g.all_engines]
+        return [e for g in self.engine_groups for e in g.all_engines]
 
     @property
     def num_new_engines(self):
-        return sum(g.num_new_engines for g in self.server_groups)
+        return sum(g.num_new_engines for g in self.engine_groups)
 
     @num_new_engines.setter
     def num_new_engines(self, value):
-        for g in self.server_groups:
+        for g in self.engine_groups:
             g.num_new_engines = value
 
     @property
     def engine_gpu_counts(self) -> list[int]:
         """Per-engine GPU count for all node-0 engines, parallel to ``engines``."""
-        return [g.num_gpus_per_engine for g in self.server_groups for _ in g.engines]
+        return [g.num_gpus_per_engine for g in self.engine_groups for _ in g.engines]
 
     @property
     def engine_gpu_offsets(self) -> list[int]:
@@ -264,7 +378,7 @@ def engine_gpu_offsets(self) -> list[int]:
         Accounts for placeholder groups that occupy GPU slots without creating engines.
         """
         offsets = []
-        for g in self.server_groups:
+        for g in self.engine_groups:
             for j in range(len(g.engines)):
                 offsets.append(g.gpu_offset + j * g.num_gpus_per_engine)
         return offsets
@@ -272,7 +386,7 @@ def engine_gpu_offsets(self) -> list[int]:
     @property
     def nodes_per_engine(self):
         """Nodes per engine.  Only valid when all active groups share the same value."""
-        values = {g.nodes_per_engine for g in self.server_groups if g.worker_type != "placeholder"}
+        values = {g.nodes_per_engine for g in self.engine_groups}
         if len(values) != 1:
             raise ValueError(f"Heterogeneous nodes_per_engine across groups: {values}")
         return values.pop()
@@ -280,12 +394,12 @@ def nodes_per_engine(self):
     def recover(self):
         """Recover dead engines across all active groups, overlapping init."""
         # Record dead indices per group before starting.
-        dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in self.server_groups]
+        dead_per_group = [[i for i, engine in enumerate(g.all_engines) if engine is None] for g in self.engine_groups]
 
         # Start all groups concurrently.
         all_handles = []
         port_cursors: dict[int, int] = {}
-        for g in self.server_groups:
+        for g in self.engine_groups:
             handles, port_cursors = g.start_engines(port_cursors)
             all_handles.extend(handles)
         if all_handles:
@@ -293,69 +407,35 @@ def recover(self):
 
         # Post-recovery: offload then onload weights for newly created engines.
         release_handles = []
-        updatable_new_engines = []
-        non_updatable_groups_engines: list[tuple[str, list]] = []
-        for g, dead_indices in zip(self.server_groups, dead_per_group, strict=True):
+        new_engines_all = []
+        for g, dead_indices in zip(self.engine_groups, dead_per_group, strict=True):
             logger.info(f"Recovered {g.num_new_engines} dead rollout engines (worker_type={g.worker_type})")
             assert g.num_new_engines == len(dead_indices), "num_new_engines does not match dead_indices length"
-            if g.needs_offload and dead_indices:
+            if g.args.offload_rollout and dead_indices:
                 new_engines = [g.all_engines[i] for i in dead_indices]
                 release_handles.extend(engine.release_memory_occupation.remote() for engine in new_engines)
-                if self.update_weights:
-                    updatable_new_engines.extend(new_engines)
-                elif g.model_path:
-                    non_updatable_groups_engines.append((g.model_path, new_engines))
+                new_engines_all.extend(new_engines)
 
         if release_handles:
             ray.get(release_handles)
-            # Resume GPU memory for all engines that need offload.
-            all_resume_engines = updatable_new_engines[:]
-            for _model_path, engines in non_updatable_groups_engines:
-                all_resume_engines.extend(engines)
-            if all_resume_engines:
-                ray.get(
-                    [
-                        engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS])
-                        for engine in all_resume_engines
-                    ]
-                )
+            ray.get(
+                [engine.resume_memory_occupation.remote(tags=[GPU_MEMORY_TYPE_WEIGHTS]) for engine in new_engines_all]
+            )
 
     def offload(self):
         """Release memory occupation across all groups (concurrent)."""
         handles = []
-        for g in self.server_groups:
+        for g in self.engine_groups:
             handles.extend(g.offload())
         return ray.get(handles) if handles else []
 
     def onload(self, tags: list[str] | None = None):
         """Resume memory occupation across all groups (concurrent)."""
         handles = []
-        for g in self.server_groups:
+        for g in self.engine_groups:
             handles.extend(g.onload(tags))
         return ray.get(handles) if handles else []
 
-    def onload_weights(self):
-        """Restore weights for offloaded groups.
-
-        All groups resume from CPU cache via ``resume_memory_occupation``.
-        For updatable servers, weights will be overwritten by
-        ``update_weights`` shortly after.  For non-updatable servers the
-        CPU backup already contains the correct (unchanged) weights.
-        """
-        handles = []
-        for g in self.server_groups:
-            if not g.needs_offload:
-                continue
-            handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS]))
-        return ray.get(handles) if handles else []
-
-    def onload_kv(self):
-        """Resume KV cache and CUDA graphs for offloaded groups."""
-        handles = []
-        for g in self.server_groups:
-            handles.extend(g.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH]))
-        return ray.get(handles) if handles else []
-
 
 @ray.remote
 class RolloutManager:
@@ -366,7 +446,8 @@ def __init__(self, args, pg):
 
         self.pg = pg
         self.args = args
-        self.rollout_backend = getattr(args, "rollout_backend", "sglang")
+
+        init_tracking(args, primary=False)
 
         data_source_cls = load_function(self.args.data_source_path)
         self.data_source = data_source_cls(args)
@@ -389,37 +470,18 @@ def __init__(self, args, pg):
         else:
             init_http_client(args)
             self.servers = start_rollout_servers(args, pg)
-
-        init_tracking(args, primary=False)
         self.rollout_engine_lock = Lock.options(num_cpus=1, num_gpus=0).remote()
         self.rollout_id = -1
 
         self._health_monitors = []
         if not self.args.debug_train_only and self.args.use_fault_tolerance:
             for srv in self.servers.values():
-                for group in srv.server_groups:
+                for group in srv.engine_groups:
                     monitor = RolloutHealthMonitor(group, args)
                     monitor.start()
                     self._health_monitors.append(monitor)
             self._ci_fault_injection_pending = self.args.ci_test  # Flag for CI fault injection
 
-    def _get_metrics_router_addr(self) -> str | None:
-        """Return the router address for scraping SGLang engine metrics.
-
-        The sglang_router gateway exposes ``/engine_metrics`` on its main port,
-        which aggregates Prometheus metrics from all backend sglang servers.
-        Returns ``http://{ip}:{port}`` for the first server, or ``None`` when
-        metrics are disabled or no servers are running.
-        """
-        srv = self.server
-        if srv is None or srv.router_ip is None:
-            return None
-        return f"http://{srv.router_ip}:{srv.router_port}"
-
-    def get_metrics_router_addr(self) -> str | None:
-        """Public wrapper for remote calls from the driver process."""
-        return self._get_metrics_router_addr()
-
     def _try_ci_fault_injection(self):
         """Try to inject fault during generate (when health monitor is running)."""
         if not self._ci_fault_injection_pending:
@@ -428,11 +490,11 @@ def _try_ci_fault_injection(self):
         # Only inject fault once
         self._ci_fault_injection_pending = False
 
-        if self.server and self.server.server_groups[0].all_engines and self.server.server_groups[0].all_engines[0]:
+        if self.server and self.server.engine_groups[0].all_engines and self.server.engine_groups[0].all_engines[0]:
             logger.info("CI Fault Injection: Simulating crash on engine 0 during generate")
             try:
                 # This will cause the ray actor to exit
-                self.server.server_groups[0].all_engines[0].simulate_crash.remote()
+                self.server.engine_groups[0].all_engines[0].simulate_crash.remote()
                 # Wait for health monitor to detect the crash and mark engine as None
                 # health_check_interval + health_check_timeout + buffer
                 wait_time = self.args.rollout_health_check_interval + self.args.rollout_health_check_timeout + 5
@@ -444,7 +506,6 @@ def _try_ci_fault_injection(self):
     def dispose(self):
         for monitor in self._health_monitors:
             monitor.stop()
-        logging_utils.finish_tracking(self.args)
 
     @property
     def server(self) -> RolloutServer | None:
@@ -453,30 +514,18 @@ def server(self) -> RolloutServer | None:
             return None
         return next(iter(self.servers.values()))
 
-    def _get_updatable_server(self) -> RolloutServer | None:
-        """Return the server with ``update_weights=True``.
-
-        When multiple updatable servers exist, returns the first one
-        (multi-model weight update is not yet supported).
-        """
-        for srv in self.servers.values():
-            if srv.update_weights:
-                return srv
-        return None
+    def _get_server(self, model_name: str | None = None) -> RolloutServer | None:
+        if model_name is None:
+            return self.server
+        return self.servers.get(model_name)
 
     @property
     def rollout_engines(self):
         """All node-0 engines across all servers / models."""
         return [e for srv in self.servers.values() for e in srv.engines]
 
-    def get_updatable_engines_and_lock(self):
-        """Return engines eligible for weight updates.
-
-        Returns engines from the first model that has
-        ``update_weights=True``.  Frozen models (reference, reward,
-        etc.) are automatically excluded.
-        """
-        srv = self._get_updatable_server()
+    def get_rollout_engines_and_lock(self, model_name: str | None = None):
+        srv = self._get_server(model_name)
         engines = srv.engines if srv else []
         gpu_counts = srv.engine_gpu_counts if srv else []
         gpu_offsets = srv.engine_gpu_offsets if srv else []
@@ -529,31 +578,15 @@ def onload(self, tags: list[str] | None = None):
             srv.onload(tags)
 
     def onload_weights(self):
-        # vLLM does not expose a separate weights-only resume: a single
-        # resume_memory_occupation reloads both weights and KV cache.  To keep
-        # the public ``onload_weights`` / ``onload_kv`` contract on the train
-        # side unchanged, we load everything here and make ``onload_kv`` a
-        # no-op for the vLLM backend.
-        if self.rollout_backend == "vllm":
-            self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS, GPU_MEMORY_TYPE_KV_CACHE])
-        else:
-            self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])
+        self.onload(tags=[GPU_MEMORY_TYPE_WEIGHTS])
 
     def onload_kv(self):
-        # See ``onload_weights``: for vLLM, weights+KV are already resumed
-        # together, so this is a no-op to avoid a double-resume.
-        if self.rollout_backend == "vllm":
-            return
         self.onload(tags=[GPU_MEMORY_TYPE_KV_CACHE, GPU_MEMORY_TYPE_CUDA_GRAPH])
 
-    def recover_updatable_engines(self):
-        """Restart any dead rollout engines and update num_new_engines for update_weights detection.
-
-        Recovers the updatable model (the one that receives weight
-        updates from training).
-        """
+    def recover_rollout_engines(self, model_name: str | None = None):
+        """Restart any dead rollout engines and update num_new_engines for update_weights detection."""
         self.health_monitoring_pause()
-        srv = self._get_updatable_server()
+        srv = self._get_server(model_name)
         if self.rollout_id == -1 or srv is None:
             engines = srv.engines if srv else []
             gpu_counts = srv.engine_gpu_counts if srv else []
@@ -569,9 +602,9 @@ def recover_updatable_engines(self):
             srv.engine_gpu_offsets,
         )
 
-    def clear_updatable_num_new_engines(self):
+    def clear_num_new_engines(self, model_name: str | None = None):
         # when fault tolerance is not enabled, we need to manually clear num_new_engines after update_weights
-        srv = self._get_updatable_server()
+        srv = self._get_server(model_name)
         if srv:
             srv.num_new_engines = 0
 
@@ -740,13 +773,9 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl
             loss_masks.append(sample.loss_mask)
         train_data["loss_masks"] = loss_masks
 
-        # Overwrite raw_reward when available. Mixed-source batches may only
-        # populate this field for a subset of samples (e.g. SWE but not code).
-        if any(sample.metadata and "raw_reward" in sample.metadata for sample in samples):
-            train_data["raw_reward"] = [
-                sample.metadata["raw_reward"] if sample.metadata and "raw_reward" in sample.metadata else sample.reward
-                for sample in samples
-            ]
+        # overwriting the raw reward
+        if samples[0].metadata and "raw_reward" in samples[0].metadata:
+            train_data["raw_reward"] = [sample.metadata["raw_reward"] for sample in samples]
 
         # For rollout buffer
         if samples[0].metadata and "round_number" in samples[0].metadata:
@@ -860,7 +889,7 @@ def _allocate_rollout_engine_addr_and_ports_normal(
     num_engines_per_node = max(1, args.num_gpus_per_node // _gpus_per_engine)
     addr_and_ports: dict[int, dict] = {}
 
-    # Track per-node port cursors so that different server groups (called
+    # Track per-node port cursors so that different engine groups (called
     # sequentially) never race for the same ports on a given node.
     node_port_cursor: dict[int, int] = {}
 
@@ -937,14 +966,14 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
     ``force_new`` is False, skip launching and return the existing values.
     When ``force_new`` is True (multi-model), always allocate a fresh port.
     """
-    if not force_new and getattr(args, "sglang_router_ip", None) is not None:
+    if not force_new and args.sglang_router_ip is not None:
         return args.sglang_router_ip, args.sglang_router_port
 
     router_ip = _wrap_ipv6(get_host_info()[1])
     if force_new:
         router_port = find_available_port(random.randint(3000, 4000))
     else:
-        router_port = getattr(args, "sglang_router_port", None)
+        router_port = args.sglang_router_port
         if router_port is None:
             router_port = find_available_port(random.randint(3000, 4000))
 
@@ -959,7 +988,7 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
         router_args.sglang_router_port = router_port
 
     else:
-        from sglang_router.launch_router import RouterArgs  # noqa: delayed import — only needed for sglang router
+        from sglang_router.launch_router import RouterArgs
 
         from slime.utils.http_utils import run_router
 
@@ -991,121 +1020,11 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
     return router_ip, router_port
 
 
-def _start_vllm_rollout_servers(args, pg) -> dict[str, RolloutServer]:
-    """Start vLLM rollout server(s) with optional SlimeRouter.
-
-    When ``args.use_slime_router`` is True, a SlimeRouter is launched first
-    and each vLLM engine registers its translation sidecar with the router.
-    Otherwise, a single engine is started and used directly (no router).
-    """
-    pg_obj, reordered_bundle_indices, reordered_gpu_ids = pg
-    tp = args.rollout_num_gpus_per_engine
-    num_gpu_per_engine_local = min(tp, args.num_gpus_per_node)
-    num_engines = max(1, args.rollout_num_gpus // num_gpu_per_engine_local)
-
-    use_router = getattr(args, "use_slime_router", False)
-
-    # --- Launch the SlimeRouter if requested ---
-    if use_router:
-        router_ip, router_port = _start_router(args, has_pd_disaggregation=False)
-        args.sglang_router_ip = router_ip
-        args.sglang_router_port = router_port
-    else:
-        router_ip = None
-        router_port = None
-
-    env_vars = {name: "1" for name in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST}
-
-    from slime.backends.vllm_utils.vllm_engine import VLLMEngine
-
-    VLLMRayActor = ray.remote(VLLMEngine)
-
-    engines = []
-    init_handles = []
-    for i in range(num_engines):
-        gpu_index = i * num_gpu_per_engine_local
-        gpu_ids = [int(reordered_gpu_ids[gpu_index + j]) for j in range(num_gpu_per_engine_local)]
-        scheduling_strategy = PlacementGroupSchedulingStrategy(
-            placement_group=pg_obj,
-            placement_group_capture_child_tasks=True,
-            placement_group_bundle_index=reordered_bundle_indices[gpu_index],
-        )
-
-        engine = VLLMRayActor.options(
-            num_cpus=0.2,
-            num_gpus=0.2,
-            scheduling_strategy=scheduling_strategy,
-            runtime_env={"env_vars": env_vars},
-        ).remote(args, rank=i, base_gpu_id=gpu_ids[0], gpu_ids=gpu_ids)
-
-        host, port = ray.get(engine._get_current_node_ip_and_free_port.remote(start_port=15000 + i * 10))
-
-        init_handles.append(
-            engine.init.remote(
-                port=port,
-                host=host,
-                router_ip=router_ip,
-                router_port=router_port,
-            )
-        )
-        engines.append(engine)
-
-        # Use the first engine's host for args if no router
-        if i == 0 and not use_router:
-            args.vllm_base_url = f"http://{host}:{port}"
-            args.sglang_router_ip = host
-            args.sglang_router_port = port
-
-    # Wait for all engines to be healthy + registered with router
-    ray.get(init_handles)
-
-    first_host = args.sglang_router_ip
-    first_port = args.sglang_router_port
-
-    group = ServerGroup(
-        args=args,
-        pg=pg,
-        all_engines=engines,
-        num_gpus_per_engine=args.rollout_num_gpus_per_engine,
-        num_new_engines=num_engines,
-        worker_type="regular",
-        rank_offset=0,
-        gpu_offset=0,
-        sglang_overrides={},
-        needs_offload=getattr(args, "offload_rollout", False),
-        model_path=getattr(args, "hf_checkpoint", None),
-        router_ip=first_host,
-        router_port=first_port,
-    )
-    return {
-        "default": RolloutServer(
-            server_groups=[group],
-            router_ip=first_host,
-            router_port=first_port,
-            model_name="default",
-        )
-    }
-def _compute_rollout_offset(args) -> int:
-    """Offset (in PG bundle slots) where rollout GPUs start."""
-    if args.debug_train_only or args.debug_rollout_only or args.colocate:
-        return 0
-    offset = args.actor_num_nodes * args.actor_num_gpus_per_node
-    return offset
-
-
-def _compute_megatron_num_gpus(args) -> int:
-    """Total number of megatron (actor + critic) GPU slots in the placement group."""
-    if args.debug_rollout_only:
-        return 0
-    num = args.actor_num_nodes * args.actor_num_gpus_per_node
-    return num
-
-
 def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
     """Start rollout servers: one per model, each with its own router.
 
     Each model defined in the sglang config gets its own router and set
-    of server groups.  Server groups within a model may have different
+    of engine groups.  Engine groups within a model may have different
     ``num_gpus_per_engine`` (e.g. for PD disaggregation where prefill
     and decode use different TP sizes).
 
@@ -1114,19 +1033,12 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
     Note: ``init_http_client`` should be called separately before this,
     as the HTTP client is shared across all servers.
     """
-    if getattr(args, "rollout_backend", "sglang") == "vllm":
-        return _start_vllm_rollout_servers(args, pg)
-
     config = _resolve_sglang_config(args)
 
     servers: dict[str, RolloutServer] = {}
     gpu_offset = 0
     engine_offset = 0
 
-    # Compute megatron GPU range for per-group offload decisions.
-    rollout_pg_offset = _compute_rollout_offset(args)
-    megatron_num_gpus = _compute_megatron_num_gpus(args)
-
     for model_idx, model_cfg in enumerate(config.models):
         model_cfg.resolve(args)
 
@@ -1138,31 +1050,16 @@ def start_rollout_servers(args, pg) -> dict[str, RolloutServer]:
             args.sglang_router_ip = router_ip
             args.sglang_router_port = router_port
 
-        server_groups: list[ServerGroup] = []
+        engine_groups: list[EngineGroup] = []
+        all_init_handles: list = []
         port_cursors: dict[int, int] = {}
 
-        has_epd = model_cfg.has_encoder_disaggregation
-
-        def _make_group(group_cfg, router_ip, router_port, overrides_extra=None):
-            nonlocal engine_offset, gpu_offset
+        for group_cfg in model_cfg.engine_groups:
             gpus_per_engine = group_cfg.num_gpus_per_engine
             num_gpu_per_engine_local = min(gpus_per_engine, args.num_gpus_per_node)
             num_engines = group_cfg.num_gpus // num_gpu_per_engine_local
 
-            group_abs_start = rollout_pg_offset + gpu_offset
-            needs_offload = args.offload_rollout and group_abs_start < megatron_num_gpus
-            overrides = dict(group_cfg.overrides)
-            if overrides_extra:
-                for k, v in overrides_extra.items():
-                    overrides.setdefault(k, v)
-            if args.offload_rollout and not needs_offload:
-                overrides.setdefault("enable_memory_saver", False)
-            logger.info(
-                f"Engine group '{group_cfg.worker_type}' gpu_offset={gpu_offset} "
-                f"(abs={group_abs_start}): needs_offload={needs_offload}"
-            )
-
-            group = ServerGroup(
+            group = EngineGroup(
                 args=args,
                 pg=pg,
                 all_engines=[None] * num_engines if group_cfg.worker_type != "placeholder" else [],
@@ -1171,73 +1068,27 @@ def _make_group(group_cfg, router_ip, router_port, overrides_extra=None):
                 worker_type=group_cfg.worker_type,
                 rank_offset=engine_offset,
                 gpu_offset=gpu_offset,
-                sglang_overrides=overrides,
-                needs_offload=needs_offload,
-                model_path=overrides.get("model_path", args.hf_checkpoint),
+                sglang_overrides=group_cfg.overrides,
                 router_ip=router_ip,
                 router_port=router_port,
             )
+            handles, port_cursors = group.start_engines(port_cursors)
+            all_init_handles.extend(handles)
+            engine_groups.append(group)
+
             engine_offset += num_engines
             gpu_offset += group_cfg.num_gpus
-            return group
-
-        if has_epd:
-            # --- Phase 1: start encoder groups, wait, collect URLs ---
-            encoder_urls: list[str] = []
-            for group_cfg in model_cfg.server_groups:
-                if group_cfg.worker_type != "encoder":
-                    continue
-                group = _make_group(group_cfg, router_ip, router_port)
-                handles, port_cursors = group.start_engines(port_cursors)
-                if handles:
-                    ray.get(handles)
-                urls = ray.get([e.get_url.remote() for e in group.engines])
-                encoder_urls.extend(u for u in urls if u is not None)
-                server_groups.append(group)
-
-            logger.info(f"EPD phase 1 done: collected {len(encoder_urls)} encoder URLs: {encoder_urls}")
-
-            # --- Phase 2: start non-encoder groups, injecting encoder URLs into
-            # language-only LLM workers. Prefill groups use this for full EPD,
-            # while regular groups allow encoder/LLM split without PD.
-            non_encoder_handles: list = []
-            for group_cfg in model_cfg.server_groups:
-                if group_cfg.worker_type == "encoder":
-                    continue
-                overrides_extra = {}
-                if encoder_urls and group_cfg.worker_type in ("prefill", "regular"):
-                    overrides_extra["language_only"] = True
-                    overrides_extra["encoder_urls"] = encoder_urls
-                group = _make_group(group_cfg, router_ip, router_port, overrides_extra=overrides_extra)
-                handles, port_cursors = group.start_engines(port_cursors)
-                non_encoder_handles.extend(handles)
-                server_groups.append(group)
-
-            if non_encoder_handles:
-                ray.get(non_encoder_handles)
-        else:
-            # No EPD — start all groups in one pass (original path).
-            all_init_handles: list = []
-            for group_cfg in model_cfg.server_groups:
-                group = _make_group(group_cfg, router_ip, router_port)
-                handles, port_cursors = group.start_engines(port_cursors)
-                all_init_handles.extend(handles)
-                server_groups.append(group)
 
-            if all_init_handles:
-                ray.get(all_init_handles)
+        if all_init_handles:
+            ray.get(all_init_handles)
 
         servers[model_cfg.name] = RolloutServer(
-            server_groups=server_groups,
+            engine_groups=engine_groups,
             router_ip=router_ip,
             router_port=router_port,
             model_name=model_cfg.name,
-            update_weights=model_cfg.update_weights,
         )
 
-    # Expose per-model router info for custom rollout functions.
-    args.sglang_model_routers = {name: (srv.router_ip, srv.router_port) for name, srv in servers.items()}
-
     return servers
 
 
@@ -1259,7 +1110,7 @@ def _resolve_sglang_config(args) -> SglangConfig:
         models=[
             ModelConfig(
                 name="default",
-                server_groups=[ServerGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)],
+                engine_groups=[EngineGroupConfig(worker_type="regular", num_gpus=args.rollout_num_gpus)],
             )
         ]
     )
@@ -1322,8 +1173,6 @@ def compute_metrics_from_samples(args, samples):
     log_dict = {}
     log_dict |= dict_add_prefix(compute_statistics(response_lengths), "response_len/")
     log_dict |= _compute_zero_std_metrics(args, samples)
-    log_dict |= _compute_spec_metrics(args, samples)
-    log_dict |= _compute_prefix_cache_metrics(args, samples)
     log_dict |= _compute_reward_cat_metrics(args, samples)
     log_dict["repetition_frac"] = np.mean([int(has_repetition(s.response)) for s in samples]).item()
     log_dict["truncated_ratio"] = np.mean([int(s.status == Sample.Status.TRUNCATED) for s in samples]).item()
diff --git a/slime/ray/train_actor.py b/slime/ray/train_actor.py
index a8ba6ddc64..87400ce130 100644
--- a/slime/ray/train_actor.py
+++ b/slime/ray/train_actor.py
@@ -58,7 +58,12 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False):
         local_rank = int(os.environ.get("LOCAL_RANK", 0))
         torch.cuda.set_device(f"cuda:{local_rank}")
 
+        # Use hybrid backend when FSDP CPU offload is enabled with a CPU backend
         backend = args.distributed_backend
+        if getattr(args, "fsdp_cpu_offload", False) and getattr(args, "fsdp_cpu_backend", None):
+            cpu_backend = args.fsdp_cpu_backend
+            backend = f"cpu:{cpu_backend},cuda:{args.distributed_backend}"
+            logger.info(f"FSDP CPU offload enabled, using hybrid backend: {backend}")
 
         dist.init_process_group(
             backend=backend,
@@ -92,8 +97,6 @@ def init(self, args, role, with_ref=False, with_opd_teacher=False):
             logger.info(f"Warning: Failed to set NUMA affinity: {e}")
 
     def clear_memory(self):
-        if self.args.debug_rollout_only:
-            return
         print_memory("before TrainRayActor.clear_memory")
         clear_memory()
         print_memory("after TrainRayActor.clear_memory")
@@ -107,7 +110,7 @@ def wake_up(self, tags):
         raise NotImplementedError
 
     @abc.abstractmethod
-    def train(self, rollout_id, rollout_data_ref, external_data=None):
+    def train(self, rollout_id, rollout_data_ref):
         raise NotImplementedError
 
     @abc.abstractmethod
@@ -118,6 +121,10 @@ def save_model(self, rollout_id, force_sync=False):
     def update_weights(self):
         raise NotImplementedError
 
+    @abc.abstractmethod
+    def connect_actor_critic(self, critic_group):
+        raise NotImplementedError
+
     @abc.abstractmethod
     def _get_parallel_config(self):
         raise NotImplementedError
diff --git a/slime/rollout/backends/__init__.py b/slime/rollout/backends/__init__.py
deleted file mode 100644
index ea9ccd8a9c..0000000000
--- a/slime/rollout/backends/__init__.py
+++ /dev/null
@@ -1,18 +0,0 @@
-from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient
-from slime.rollout.backends.vllm_client import VLLMClient
-
-
-def __getattr__(name):
-    if name == "SGLangClient":
-        from slime.rollout.backends.sglang_client import SGLangClient
-
-        return SGLangClient
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-
-__all__ = [
-    "BackendCapabilities",
-    "RolloutBackendClient",
-    "SGLangClient",
-    "VLLMClient",
-]
diff --git a/slime/rollout/backends/base_client.py b/slime/rollout/backends/base_client.py
deleted file mode 100644
index 1730d6cbac..0000000000
--- a/slime/rollout/backends/base_client.py
+++ /dev/null
@@ -1,31 +0,0 @@
-from abc import ABC, abstractmethod
-from dataclasses import dataclass
-
-from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse
-
-
-@dataclass
-class BackendCapabilities:
-    supports_abort: bool
-    supports_routed_experts: bool
-    supports_prompt_logprobs: bool
-
-
-class RolloutBackendClient(ABC):
-    @property
-    @abstractmethod
-    def capabilities(self) -> BackendCapabilities:
-        ...
-
-    @abstractmethod
-    async def generate(
-        self,
-        request: RolloutBackendRequest,
-        base_url: str,
-        headers: dict | None = None,
-    ) -> RolloutBackendResponse:
-        ...
-
-    async def abort(self) -> list[str]:
-        """Return worker URLs for abort. Empty for backends without worker-level abort."""
-        return []
diff --git a/slime/rollout/backends/sglang_client.py b/slime/rollout/backends/sglang_client.py
deleted file mode 100644
index 2aae5f80ee..0000000000
--- a/slime/rollout/backends/sglang_client.py
+++ /dev/null
@@ -1,88 +0,0 @@
-import logging
-
-import numpy as np
-import pybase64
-import sglang_router
-from packaging.version import parse
-
-from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient
-from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse
-from slime.utils.http_utils import get, post
-
-logger = logging.getLogger(__name__)
-
-
-class SGLangClient(RolloutBackendClient):
-    def __init__(self, args):
-        self.args = args
-
-    @property
-    def capabilities(self) -> BackendCapabilities:
-        return BackendCapabilities(
-            supports_abort=True,
-            supports_routed_experts=bool(getattr(self.args, "use_rollout_routing_replay", False)),
-            supports_prompt_logprobs=True,
-        )
-
-    async def generate(
-        self,
-        request: RolloutBackendRequest,
-        base_url: str,
-        headers: dict | None = None,
-    ) -> RolloutBackendResponse:
-        payload = {
-            "sampling_params": request.sampling_params,
-            "return_logprob": request.return_logprob,
-            "return_routed_experts": request.return_routed_experts,
-        }
-        # For multimodal: send raw text so SGLang's server-side processor can
-        # expand image placeholders. Otherwise send pre-tokenized input_ids.
-        if request.image_data and request.text is not None:
-            payload["text"] = request.text
-            payload["image_data"] = request.image_data
-        else:
-            payload["input_ids"] = request.input_ids
-            if request.image_data:
-                payload["image_data"] = request.image_data
-
-        url = f"{base_url.rstrip('/')}/generate"
-        output = await post(url, payload, headers=headers)
-
-        meta = output.get("meta_info", {})
-        logprobs = meta.get("output_token_logprobs", [])
-        output_token_ids = [item[1] for item in logprobs]
-        output_token_logprobs = [item[0] for item in logprobs]
-
-        finish_reason = meta.get("finish_reason", {}).get("type", "stop")
-        routed_experts = None
-        if "routed_experts" in meta and self.capabilities.supports_routed_experts:
-            num_layers = getattr(self.args, "num_layers", 0)
-            moe_topk = getattr(self.args, "moe_router_topk", 1)
-            if num_layers and moe_topk:
-                routed_experts = np.frombuffer(
-                    pybase64.b64decode(meta["routed_experts"].encode("ascii")),
-                    dtype=np.int32,
-                ).reshape(
-                    len(request.input_ids) + len(output_token_ids) - 1,
-                    num_layers,
-                    moe_topk,
-                )
-
-        return RolloutBackendResponse(
-            text=output.get("text", ""),
-            output_token_ids=output_token_ids,
-            output_token_logprobs=output_token_logprobs,
-            finish_reason=finish_reason,
-            prompt_tokens=meta.get("prompt_tokens", len(request.input_ids)),
-            completion_tokens=meta.get("completion_tokens", len(output_token_ids)),
-            backend_raw=output,
-            routed_experts=routed_experts,
-        )
-
-    async def abort(self) -> list[str]:
-        base = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}"
-        if parse(sglang_router.__version__) <= parse("0.2.1") or getattr(self.args, "use_slime_router", False):
-            r = await get(f"{base}/list_workers")
-            return r.get("urls", [])
-        r = await get(f"{base}/workers")
-        return [w["url"] for w in r.get("workers", [])]
diff --git a/slime/rollout/backends/vllm_client.py b/slime/rollout/backends/vllm_client.py
deleted file mode 100644
index 56d57d8954..0000000000
--- a/slime/rollout/backends/vllm_client.py
+++ /dev/null
@@ -1,192 +0,0 @@
-import logging
-
-from slime.rollout.backends.base_client import BackendCapabilities, RolloutBackendClient
-from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse
-from slime.utils.http_utils import get, post
-
-logger = logging.getLogger(__name__)
-
-_FINISH_REASON_MAP = {
-    "stop": "stop",
-    "length": "length",
-    "abort": "abort",
-    "end_turn": "stop",
-    "max_tokens": "length",
-}
-
-
-class VLLMClient(RolloutBackendClient):
-    """Rollout backend client for vLLM.
-
-    Supports two modes:
-
-    1. **Direct mode** (default): sends requests directly to vLLM's
-       ``/v1/completions`` endpoint and parses the native vLLM response.
-    2. **Router mode** (``use_slime_router=True``): sends SGLang-format
-       requests to ``/generate`` through the SlimeRouter, which forwards
-       to the translation sidecar. The sidecar handles the vLLM translation
-       and returns SGLang-format responses.
-    """
-
-    def __init__(self, args):
-        self.args = args
-        self._max_retries = getattr(args, "vllm_max_retries", 3)
-        self._use_router = getattr(args, "use_slime_router", False)
-
-    @property
-    def capabilities(self) -> BackendCapabilities:
-        return BackendCapabilities(
-            supports_abort=self._use_router,  # abort is supported through the sidecar
-            supports_routed_experts=False,
-            supports_prompt_logprobs=False,
-        )
-
-    async def generate(
-        self,
-        request: RolloutBackendRequest,
-        base_url: str,
-        headers: dict | None = None,
-    ) -> RolloutBackendResponse:
-        if self._use_router:
-            return await self._generate_via_router(request, base_url, headers)
-        return await self._generate_direct(request, base_url, headers)
-
-    # ------------------------------------------------------------------
-    # Router mode: SGLang-format /generate → sidecar → vLLM
-    # ------------------------------------------------------------------
-
-    async def _generate_via_router(
-        self,
-        request: RolloutBackendRequest,
-        base_url: str,
-        headers: dict | None = None,
-    ) -> RolloutBackendResponse:
-        """Send SGLang-format request through the SlimeRouter → sidecar pipeline."""
-
-        payload = {
-            "sampling_params": request.sampling_params,
-            "return_logprob": request.return_logprob,
-            "stream": False,
-        }
-        # Multimodal: send raw text + image_data so the translation sidecar can
-        # apply the chat template and forward to vLLM. Otherwise send tokenized
-        # input_ids.
-        if request.image_data and request.text is not None:
-            payload["text"] = request.text
-            payload["image_data"] = request.image_data
-        else:
-            payload["input_ids"] = request.input_ids
-            if request.image_data:
-                payload["image_data"] = request.image_data
-
-        url = f"{base_url.rstrip('/')}/generate"
-        output = await post(url, payload, headers=headers)
-
-        # Parse SGLang-format response (produced by translation sidecar)
-        meta = output.get("meta_info", {})
-        logprobs_data = meta.get("output_token_logprobs", [])
-
-        # output_token_logprobs is list of [logprob, token_id]
-        output_token_ids = [item[1] for item in logprobs_data] if logprobs_data else []
-        output_token_logprobs = [item[0] for item in logprobs_data] if logprobs_data else []
-
-        # Fall back to output_ids if logprobs not available
-        if not output_token_ids:
-            output_token_ids = output.get("output_ids") or []
-
-        finish_reason = meta.get("finish_reason", {}).get("type", "stop")
-
-        return RolloutBackendResponse(
-            text=output.get("text", ""),
-            output_token_ids=output_token_ids,
-            output_token_logprobs=output_token_logprobs,
-            finish_reason=finish_reason,
-            prompt_tokens=meta.get("prompt_tokens", len(request.input_ids)),
-            completion_tokens=len(output_token_ids),
-            backend_raw=output,
-            routed_experts=None,
-        )
-
-    # ------------------------------------------------------------------
-    # Direct mode: vLLM /v1/completions (no router)
-    # ------------------------------------------------------------------
-
-    async def _generate_direct(
-        self,
-        request: RolloutBackendRequest,
-        base_url: str,
-        headers: dict | None = None,
-    ) -> RolloutBackendResponse:
-        sp = request.sampling_params
-        # vLLM /v1/completions accepts a string or list of token IDs as "prompt".
-        # Prefer raw text when provided (e.g. multimodal callers); fall back to
-        # the pre-tokenized input_ids otherwise.
-        prompt_field = request.text if request.text is not None else request.input_ids
-        payload = {
-            "prompt": prompt_field,
-            "max_tokens": sp.get("max_new_tokens", 1024),
-            "temperature": sp.get("temperature", 1.0),
-            "top_p": sp.get("top_p", 1.0),
-            "top_k": sp.get("top_k", -1),
-            "stop": sp.get("stop"),
-            "stop_token_ids": sp.get("stop_token_ids"),
-            "skip_special_tokens": sp.get("skip_special_tokens", True),
-            "spaces_between_special_tokens": sp.get("spaces_between_special_tokens", False),
-            "logprobs": 1 if request.return_logprob else None,
-            "include_stop_str_in_output": sp.get("no_stop_trim", False),
-            "return_token_ids": True,
-            "seed": sp.get("sampling_seed"),
-        }
-        payload = {k: v for k, v in payload.items() if v is not None}
-
-        url = f"{base_url.rstrip('/')}/v1/completions"
-        output = await post(url, payload, headers=headers)
-
-        choice = output.get("choices", [{}])[0]
-        text = choice.get("text", "")
-        finish = choice.get("finish_reason", "stop")
-        finish_reason = _FINISH_REASON_MAP.get(finish, "stop")
-
-        # vLLM /v1/completions response format:
-        #   choice.token_ids: list[int]  (output token IDs)
-        #   choice.logprobs.token_logprobs: list[float|None]
-        #   choice.logprobs.tokens: list[str]  (token text, not IDs)
-        output_token_ids = choice.get("token_ids") or []
-        logprobs_obj = choice.get("logprobs") or {}
-        raw_logprobs = logprobs_obj.get("token_logprobs") or []
-        output_token_logprobs = [float(lp) if lp is not None else 0.0 for lp in raw_logprobs]
-
-        # If token_ids not in response, fall back to tokenizer
-        if not output_token_ids and text:
-            logger.warning("vLLM response missing token_ids, falling back to tokenizer")
-            from slime.utils.processing_utils import load_tokenizer
-            tokenizer = load_tokenizer(self.args.hf_checkpoint, trust_remote_code=True)
-            output_token_ids = tokenizer.encode(text, add_special_tokens=False)
-
-        # Ensure logprobs list matches token count
-        if len(output_token_logprobs) < len(output_token_ids):
-            output_token_logprobs.extend([0.0] * (len(output_token_ids) - len(output_token_logprobs)))
-
-        usage = output.get("usage", {})
-        return RolloutBackendResponse(
-            text=text,
-            output_token_ids=output_token_ids,
-            output_token_logprobs=output_token_logprobs,
-            finish_reason=finish_reason,
-            prompt_tokens=usage.get("prompt_tokens", len(request.input_ids)),
-            completion_tokens=usage.get("completion_tokens", len(output_token_ids)),
-            backend_raw=output,
-            routed_experts=None,
-        )
-
-    # ------------------------------------------------------------------
-    # Abort (router mode only)
-    # ------------------------------------------------------------------
-
-    async def abort(self) -> list[str]:
-        """Return worker URLs for abort. Only works in router mode."""
-        if not self._use_router:
-            return []
-        base = f"http://{self.args.sglang_router_ip}:{self.args.sglang_router_port}"
-        r = await get(f"{base}/list_workers")
-        return r.get("urls", [])
diff --git a/slime/rollout/base_types.py b/slime/rollout/base_types.py
index aefa73b15d..f0eb0d96ce 100644
--- a/slime/rollout/base_types.py
+++ b/slime/rollout/base_types.py
@@ -4,39 +4,6 @@
 from slime.utils.types import Sample
 
 
-@dataclass
-class RolloutBackendRequest:
-    """Backend-agnostic rollout request.
-
-    For multimodal requests, callers should set both ``text`` (the raw prompt
-    with image placeholders) and ``image_data``. Backends that natively expand
-    image placeholders (SGLang) will prefer ``text`` over ``input_ids`` in this
-    case so the server-side processor can do the expansion.
-    """
-
-    input_ids: list[int]
-    sampling_params: dict[str, Any]
-    return_logprob: bool = True
-    return_routed_experts: bool = False
-    image_data: list[str] | None = None
-    session_id: str | None = None
-    text: str | None = None
-
-
-@dataclass
-class RolloutBackendResponse:
-    """Backend-agnostic rollout response."""
-
-    text: str
-    output_token_ids: list[int]
-    output_token_logprobs: list[float]
-    finish_reason: str  # "stop" | "length" | "abort"
-    prompt_tokens: int
-    completion_tokens: int
-    backend_raw: dict
-    routed_experts: Any = None
-
-
 @dataclass
 class RolloutFnTrainOutput:
     samples: list[list[Sample]]
diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py
index ca7171ecae..e52cd27ef5 100644
--- a/slime/rollout/data_source.py
+++ b/slime/rollout/data_source.py
@@ -58,7 +58,7 @@ def __init__(self, args):
         # TODO remove this
         self.metadata = {}
 
-        if args.rollout_global_dataset and args.prompt_data is not None:
+        if args.rollout_global_dataset:
             tokenizer = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
             processor = load_processor(args.hf_checkpoint, trust_remote_code=True)
 
@@ -156,12 +156,10 @@ def load(self, rollout_id=None):
         self.sample_index = state_dict.get("sample_index", 0)
         self.metadata = state_dict.get("metadata", {})
 
-        if self.args.rollout_global_dataset and self.args.rollout_shuffle and self.dataset is not None:
+        if self.args.rollout_global_dataset and self.args.rollout_shuffle:
             self.dataset.shuffle(self.epoch_id)
 
     def __len__(self) -> int:
-        if self.dataset is None:
-            return 0
         return len(self.dataset)
 
 
diff --git a/slime/rollout/filter_hub/dynamic_sampling_filters.py b/slime/rollout/filter_hub/dynamic_sampling_filters.py
index 9c579d0fda..dc17422d68 100644
--- a/slime/rollout/filter_hub/dynamic_sampling_filters.py
+++ b/slime/rollout/filter_hub/dynamic_sampling_filters.py
@@ -8,7 +8,7 @@
 
 def check_reward_nonzero_std(args, samples: list[Sample], **kwargs):
     rewards = [sample.get_reward_value(args) for sample in samples]
-    keep = torch.tensor(rewards, dtype=torch.float64).std() > 1e-6
+    keep = torch.tensor(rewards, dtype=torch.float).std() > 0.0
     return DynamicFilterOutput(
         keep=keep,
         reason=None if keep else f"zero_std_{round(rewards[0], 1)}",
diff --git a/slime/rollout/generate_hub/__init__.py b/slime/rollout/generate_hub/__init__.py
new file mode 100644
index 0000000000..2d43c62459
--- /dev/null
+++ b/slime/rollout/generate_hub/__init__.py
@@ -0,0 +1 @@
+# TODO: maybe move `sglang_rollout::generate` to this folder
diff --git a/slime/rollout/generate_hub/benchmarkers.py b/slime/rollout/generate_hub/benchmarkers.py
new file mode 100644
index 0000000000..facf089185
--- /dev/null
+++ b/slime/rollout/generate_hub/benchmarkers.py
@@ -0,0 +1,25 @@
+import logging
+import random
+from argparse import Namespace
+from copy import deepcopy
+from typing import Any
+
+from slime.rollout.sglang_rollout import generate as _generate_base
+from slime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+
+async def generate_with_random_osl(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample:
+    # TODO: make it configurable after we have an enhanced arg parser
+    min_osl = 32 * 1024
+    max_osl = 64 * 1024
+
+    modified_sampling_params = deepcopy(sampling_params)
+    modified_sampling_params["ignore_eos"] = True
+    modified_sampling_params["max_new_tokens"] = random.randrange(min_osl, max_osl)
+
+    ans = await _generate_base(args, sample, modified_sampling_params)
+
+    logger.info(f"generate_with_random_osl {ans.response_length=}")
+    return ans
diff --git a/slime/rollout/on_policy_distillation.py b/slime/rollout/on_policy_distillation.py
index 9190974345..cf25b4897d 100644
--- a/slime/rollout/on_policy_distillation.py
+++ b/slime/rollout/on_policy_distillation.py
@@ -1,7 +1,6 @@
 import aiohttp
 import torch
 
-from slime.utils.processing_utils import encode_image_for_rollout_engine
 from slime.utils.types import Sample
 
 
@@ -17,11 +16,6 @@ async def reward_func(args, sample, **kwargs):
         "return_logprob": True,
         "logprob_start_len": 0,
     }
-
-    if sample.multimodal_inputs and sample.multimodal_inputs.get("images"):
-        image_data = sample.multimodal_inputs["images"]
-        payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data]
-
     session_kwargs = {}
     async with aiohttp.ClientSession(**session_kwargs) as session:
         async with session.post(args.rm_url, json=payload) as resp:
diff --git a/slime/rollout/rm_hub/gpqa.py b/slime/rollout/rm_hub/gpqa.py
index affb7a22ad..df2e3c0b91 100644
--- a/slime/rollout/rm_hub/gpqa.py
+++ b/slime/rollout/rm_hub/gpqa.py
@@ -2,7 +2,7 @@
 import string
 from collections.abc import Iterable
 
-DEFAULT_VALID_LETTERS = list(string.ascii_uppercase[:10])
+DEFAULT_VALID_LETTERS = list(string.ascii_uppercase[:8])
 
 
 def _strip_chain_of_thought(text: str) -> str:
diff --git a/slime/rollout/sft_rollout.py b/slime/rollout/sft_rollout.py
index 4407463d23..6b914a964e 100644
--- a/slime/rollout/sft_rollout.py
+++ b/slime/rollout/sft_rollout.py
@@ -47,10 +47,6 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
         tools = sample.metadata.get("tools", None)
 
         token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools)
-        if len(token_ids) != len(loss_mask):
-            raise ValueError(
-                f"SFT rollout produced mismatched token_ids/loss_mask lengths: {len(token_ids)=}, {len(loss_mask)=}"
-            )
 
         response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]
 
diff --git a/slime/rollout/sglang_rollout.py b/slime/rollout/sglang_rollout.py
index 9b6bc3cbc4..042fdab315 100644
--- a/slime/rollout/sglang_rollout.py
+++ b/slime/rollout/sglang_rollout.py
@@ -9,14 +9,17 @@
 from typing import Any
 
 import numpy as np
+import pybase64
+import sglang_router
+from packaging.version import parse
 from tqdm import tqdm
 
-from slime.rollout.base_types import RolloutBackendRequest, RolloutBackendResponse, RolloutFnEvalOutput, RolloutFnTrainOutput
+from slime.rollout.base_types import RolloutFnEvalOutput, RolloutFnTrainOutput
 from slime.rollout.filter_hub.base_types import MetricGatherer, call_dynamic_filter
 from slime.utils.async_utils import run
 from slime.utils.data import Dataset
 from slime.utils.eval_config import EvalDatasetConfig
-from slime.utils.http_utils import post
+from slime.utils.http_utils import get, post
 from slime.utils.misc import SingletonMeta, load_function
 from slime.utils.processing_utils import (
     build_processor_kwargs,
@@ -24,89 +27,14 @@
     load_processor,
     load_tokenizer,
 )
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_function, trace_span
 from slime.utils.types import Sample
 
 from .rm_hub import async_rm, batched_async_rm
 
-__all__ = ["generate_rollout", "get_model_url"]
-
-
-def _get_backend_client(args):
-    backend = getattr(args, "rollout_backend", "sglang")
-    if backend == "vllm":
-        from .backends.vllm_client import VLLMClient
-        return VLLMClient(args)
-    from .backends.sglang_client import SGLangClient
-    return SGLangClient(args)
-
-
-def _apply_backend_response(sample, resp: RolloutBackendResponse, args):
-    """Apply RolloutBackendResponse to sample."""
-    sample.tokens = sample.tokens + resp.output_token_ids
-    sample.response_length += len(resp.output_token_ids)
-    sample.response += resp.text
-    if sample.loss_mask is not None:
-        assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
-        sample.loss_mask += [1] * len(resp.output_token_ids)
-    if sample.rollout_log_probs is None:
-        sample.rollout_log_probs = []
-    sample.rollout_log_probs += resp.output_token_logprobs
-    if resp.routed_experts is not None:
-        sample.rollout_routed_experts = resp.routed_experts
-    meta = {
-        "finish_reason": {"type": resp.finish_reason},
-        "prompt_tokens": resp.prompt_tokens,
-        "completion_tokens": resp.completion_tokens,
-        **{k: v for k, v in resp.backend_raw.get("meta_info", {}).items() if k not in ("finish_reason",)},
-    }
-    sample.update_from_meta_info(args, meta)
+__all__ = ["generate_rollout"]
 
 logger = logging.getLogger(__name__)
 
-_PROCESSOR_PROMPT_KEYS = {"input_ids", "attention_mask"}
-
-
-def _prepare_prompt_ids(sample: Sample, tokenizer, processor: Any) -> list[int]:
-    raw_multimodal_inputs = sample.multimodal_inputs or {}
-    has_multimodal_inputs = any(value is not None for value in raw_multimodal_inputs.values())
-    reuse_existing_input_ids = bool(sample.tokens) and (
-        sample.multimodal_train_inputs is not None or not has_multimodal_inputs
-    )
-
-    if processor and has_multimodal_inputs and not reuse_existing_input_ids:
-        processor_output = processor(text=sample.prompt, **build_processor_kwargs(raw_multimodal_inputs))
-        prompt_ids = processor_output["input_ids"][0]
-        if sample.multimodal_train_inputs is None:
-            sample.multimodal_train_inputs = {
-                k: v for k, v in processor_output.items() if k not in _PROCESSOR_PROMPT_KEYS
-            } or None
-        return prompt_ids
-
-    if reuse_existing_input_ids:
-        return sample.tokens
-
-    return tokenizer.encode(sample.prompt, add_special_tokens=False)
-
-
-def get_model_url(args: Namespace, model_name: str, endpoint: str = "/generate") -> str:
-    """Return the router URL for a named model.
-
-    Use this in custom rollout functions to route requests to a specific
-    model when multiple models are deployed via ``--sglang-config``::
-
-        url = get_model_url(args, "ref", "/generate")
-        resp = await post(url, json=payload)
-
-    Falls back to the default router if *model_name* is not found or
-    ``sglang_model_routers`` is not set.
-    """
-    routers = getattr(args, "sglang_model_routers", None)
-    if routers and model_name in routers:
-        ip, port = routers[model_name]
-        return f"http://{ip}:{port}{endpoint}"
-    return f"http://{args.sglang_router_ip}:{args.sglang_router_port}{endpoint}"
-
 
 class GenerateState(metaclass=SingletonMeta):
     """
@@ -120,7 +48,7 @@ def __init__(self, args: Namespace) -> None:
         self.processor = load_processor(args.hf_checkpoint, trust_remote_code=True)
 
         self.semaphore = asyncio.Semaphore(
-            args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
+            args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
         )
         self.sampling_params: dict[str, Any] = dict(
             temperature=args.rollout_temperature,
@@ -178,16 +106,29 @@ def submit_generate_tasks(self, samples: list[list[Sample]]) -> None:
 
 
 async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, Any]) -> Sample:
-    """Generate using backend client (SGLang or vLLM)."""
+    """Generate using traditional SGLang router with token-based workflow"""
     if args.ci_test:
         assert isinstance(sample.prompt, str)
 
     state = GenerateState(args)
+    url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
+
     assert (
         sample.status == Sample.Status.PENDING or sample.status == Sample.Status.ABORTED
     ), f"Sample status is {sample.status}"
 
-    prompt_ids = _prepare_prompt_ids(sample, state.tokenizer, state.processor)
+    if state.processor and sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()):
+        processor_kwargs = build_processor_kwargs(sample.multimodal_inputs)
+        processor_output = state.processor(text=sample.prompt, **processor_kwargs)
+        prompt_ids = processor_output["input_ids"][0]
+        sample.multimodal_train_inputs = {
+            k: v for k, v in processor_output.items() if k not in ["input_ids", "attention_mask"]
+        } or None
+    else:
+        prompt_ids = state.tokenizer.encode(sample.prompt, add_special_tokens=False)
+
+    if len(sample.response) > 0:
+        sampling_params["max_new_tokens"] -= len(sample.tokens) - len(prompt_ids)
 
     assert (
         sampling_params["max_new_tokens"] >= 0
@@ -196,70 +137,74 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         sample.status = Sample.Status.TRUNCATED
         return sample
 
-    input_ids = sample.tokens if len(sample.response) > 0 else prompt_ids
-    if not sample.tokens:
-        sample.tokens = list(input_ids)
+    # Prepare payload for sglang server
+    payload = {
+        "sampling_params": sampling_params,
+        "return_logprob": True,
+    }
+
+    if args.use_rollout_routing_replay:
+        payload["return_routed_experts"] = True
 
-    # Multimodal images: when present, prefer sending raw text+image_data so the
-    # backend (SGLang) expands image placeholders with its own processor rules.
-    images = sample.multimodal_inputs.get("images") if sample.multimodal_inputs else None
-    image_data = [encode_image_for_rollout_engine(img) for img in images] if images else None
+    if sample.multimodal_inputs and sample.multimodal_inputs["images"]:
+        image_data = sample.multimodal_inputs["images"]
+        payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data]
+
+    # Use existing tokens for multi-turn or tokenize the new prompt
+    if len(sample.response) > 0:
+        payload["input_ids"] = sample.tokens
+    else:
+        payload["input_ids"] = prompt_ids
+        if not sample.tokens:  # Initialize sample.tokens for the first turn
+            sample.tokens = prompt_ids
 
-    use_radix = args.use_slime_router and "RadixTreeMiddleware" in getattr(args, "slime_router_middleware_paths", [])
-    if use_radix:
-        # RadixTree path: SGLang-specific, keep direct post
+    # Use session_id for consistent hashing routing if router uses consistent_hashing policy
+    headers = None
+    if args.sglang_router_policy == "consistent_hashing" and sample.session_id:
+        headers = {"X-SMG-Routing-Key": sample.session_id}
+
+    output = await post(url, payload, headers=headers)
+
+    if args.use_slime_router and "RadixTreeMiddleware" in args.slime_router_middleware_paths:
         from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree
 
-        url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
-        payload = {
-            "sampling_params": sampling_params,
-            "return_logprob": True,
-            "return_routed_experts": args.use_rollout_routing_replay,
-        }
-        if image_data:
-            payload["image_data"] = image_data
-            payload["text"] = sample.prompt
-        else:
-            payload["input_ids"] = input_ids
-        headers = (
-            {"X-SMG-Routing-Key": sample.session_id}
-            if getattr(args, "sglang_router_policy", None) == "consistent_hashing" and sample.session_id
-            else None
-        )
-        output = await post(url, payload, headers=headers)
         sample = await postprocess_sample_with_radix_tree(args, sample, output)
     else:
-        backend = _get_backend_client(args)
-        base_url = getattr(args, "vllm_base_url", None) or f"http://{args.sglang_router_ip}:{args.sglang_router_port}"
-        headers = (
-            {"X-SMG-Routing-Key": sample.session_id}
-            if getattr(args, "sglang_router_policy", None) == "consistent_hashing" and sample.session_id
-            else None
-        )
-        req = RolloutBackendRequest(
-            input_ids=input_ids,
-            text=sample.prompt if image_data else None,
-            sampling_params=sampling_params,
-            return_logprob=True,
-            return_routed_experts=args.use_rollout_routing_replay,
-            image_data=image_data,
-            session_id=sample.session_id,
+        if "output_token_logprobs" in output["meta_info"]:
+            new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]]
+            new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]]
+        else:
+            new_response_tokens, new_response_log_probs = [], []
+
+        # Update sample with tokens directly - avoiding re-tokenization
+        sample.tokens = sample.tokens + new_response_tokens
+        sample.response_length += len(new_response_tokens)
+        sample.response += output["text"]
+
+        # When partial rollout and masking off policy is enabled, update the loss mask
+        if sample.loss_mask is not None:
+            assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
+            sample.loss_mask += [1] * len(new_response_tokens)
+
+        if sample.rollout_log_probs is None:
+            sample.rollout_log_probs = []
+        sample.rollout_log_probs += new_response_log_probs
+
+    if "routed_experts" in output["meta_info"]:
+        sample.rollout_routed_experts = np.frombuffer(
+            pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")),
+            dtype=np.int32,
+        ).reshape(
+            len(sample.tokens) - 1,
+            args.num_layers,
+            args.moe_router_topk,
         )
-        with trace_span(
-            sample,
-            "rollout_generate",
-            attrs={"max_new_tokens": sampling_params["max_new_tokens"]},
-        ) as span:
-            resp = await backend.generate(req, base_url, headers=headers)
-            meta_info = resp.backend_raw.get("meta_info") if isinstance(resp.backend_raw, dict) else None
-            if meta_info:
-                span.update(build_sglang_meta_trace_attrs(meta_info))
-        _apply_backend_response(sample, resp, args)
+
+    sample.update_from_meta_info(args, output["meta_info"])
 
     return sample
 
 
-@trace_function("generate_and_rm", target="sample")
 async def generate_and_rm(
     args: Namespace,
     sample: Sample | list[Sample],
@@ -303,33 +248,28 @@ async def generate_and_rm(
     if args.group_rm:
         return sample
 
+    # multi samples
     if isinstance(sample, list):
         samples = sample
-        if any(sample.status == Sample.Status.ABORTED for sample in samples):
+        if any([sample.status == Sample.Status.ABORTED for sample in samples]):
             return samples
 
+        # for multi agent system, the reward of some sample is calculated during generation.
         samples_need_reward = [sample for sample in samples if sample.reward is None]
-        with trace_span(samples_need_reward, "reward_model"):
-            rewards = await batched_async_rm(args, samples_need_reward)
+        rewards = await batched_async_rm(args, samples_need_reward)
         for sample, reward in zip(samples_need_reward, rewards, strict=False):
             sample.reward = reward
         return samples
     else:
         if sample.status == Sample.Status.ABORTED:
             return sample
-        # Some custom generate paths may have already filled the reward.
+        # for multi-turn environment, a reward could be assigned to the agent.
         if sample.reward is None:
-            with trace_span(sample, "reward_model"):
-                sample.reward = await async_rm(args, sample)
+            sample.reward = await async_rm(args, sample)
 
     return sample
 
 
-@trace_function(
-    "generate_and_rm_group",
-    target="group",
-    attrs_getter=lambda args, group, sampling_params, evaluation=False: {"group_size": len(group)},
-)
 async def generate_and_rm_group(
     args: Namespace, group: list[Sample], sampling_params: dict[str, Any], evaluation: bool = False
 ) -> list[Sample]:
@@ -357,8 +297,7 @@ async def generate_and_rm_group(
 
     # for the rm that need the whole group, we will do the rm here
     if not state.aborted and args.group_rm:
-        with trace_span(group, "group_reward_model"):
-            rewards = await batched_async_rm(args, group)
+        rewards = await batched_async_rm(args, group)
         for sample, reward in zip(group, rewards, strict=False):
             sample.reward = reward
 
@@ -367,19 +306,24 @@ async def generate_and_rm_group(
 
 async def abort(args: Namespace, rollout_id: int) -> list[list[Sample]]:
     aborted_samples = []
+
     state = GenerateState(args)
     assert not state.aborted
     state.aborted = True
 
-    backend = _get_backend_client(args)
-    urls = await backend.abort()
-    if urls:
-        logger.info(f"Abort request for {urls}")
-        abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls]
-        abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True)
-        for url, result in zip(urls, abort_results, strict=False):
-            if isinstance(result, Exception):
-                logger.warning(f"Failed to abort worker at {url}: {result}")
+    if parse(sglang_router.__version__) <= parse("0.2.1") or args.use_slime_router:
+        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/list_workers")
+        urls = response["urls"]
+    else:
+        response = await get(f"http://{args.sglang_router_ip}:{args.sglang_router_port}/workers")
+        urls = [worker["url"] for worker in response["workers"]]
+
+    logger.info(f"Abort request for {urls}")
+    abort_tasks = [post(f"{url}/abort_request", {"abort_all": True}) for url in urls]
+    abort_results = await asyncio.gather(*abort_tasks, return_exceptions=True)
+    for url, result in zip(urls, abort_results, strict=False):
+        if isinstance(result, Exception):
+            logger.warning(f"Failed to abort worker at {url}: {result}")
 
     # make sure all the pending tasks are finished
     count = 0
@@ -591,11 +535,10 @@ async def eval_rollout_single_dataset(
     for coro in asyncio.as_completed(tasks):
         sample = await coro
         if do_print:
-            logged_sample = sample[0] if isinstance(sample, list) else sample
             logger.info(
                 "eval_rollout_single_dataset example data: "
-                f"{[str(logged_sample.prompt) + logged_sample.response]} "
-                f"reward={logged_sample.reward}"
+                f"{[str(sample.prompt) + sample.response]} "
+                f"reward={sample.reward}"
             )
             do_print = False
         if isinstance(sample, list):
@@ -637,6 +580,5 @@ def generate_rollout(
         return output
 
     output, aborted_samples = run(generate_rollout_async(args, rollout_id, data_source.get_samples))
-    if aborted_samples:
-        data_source.add_samples(aborted_samples)
+    data_source.add_samples(aborted_samples)
     return output
diff --git a/slime/router/__init__.py b/slime/router/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/slime/router/middleware_hub/__init__.py b/slime/router/middleware_hub/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/slime/router/middleware_hub/radix_tree.py b/slime/router/middleware_hub/radix_tree.py
new file mode 100644
index 0000000000..6e722f1e25
--- /dev/null
+++ b/slime/router/middleware_hub/radix_tree.py
@@ -0,0 +1,681 @@
+from __future__ import annotations
+
+"""
+String-based Radix Trie for efficient prefix matching and token caching.
+Optimized for string prefixes with corresponding token IDs.
+"""
+
+import threading
+import time
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class MatchResult:
+    """Result of prefix matching operation."""
+
+    matched_prefix: str
+    token_ids: list[int]
+    logp: list[float]
+    loss_mask: list[int]  # Added loss mask for model generation parts
+    remaining_string: str
+    last_node: StringTreeNode
+
+
+class StringTreeNode:
+    """Tree node for string-based radix trie."""
+
+    counter = 0
+
+    def __init__(self, node_id: int | None = None):
+        # Core tree structure
+        self.children: list[StringTreeNode] = []  # Use list to store children
+        self.parent: StringTreeNode | None = None
+
+        # Node data
+        self.string_key: str = ""  # The string fragment this node represents
+        self.token_ids: list[int] | None = None  # Token IDs for this node only (not cumulative)
+        self.logp: list[float] | None = None  # Log probabilities for this node's tokens
+        self.loss_mask: list[int] | None = None  # Loss mask for model generation parts
+
+        # Access tracking
+        self.last_access_time = time.monotonic()
+        self.access_count = 0
+
+        # Reference counting for protection from eviction
+        self.ref_count = 0
+
+        # Weight version tracking
+        self.weight_version: int | None = None  # Weight version for this node
+
+        # Node identification
+        self.id = StringTreeNode.counter if node_id is None else node_id
+        StringTreeNode.counter += 1
+
+    @property
+    def is_leaf(self) -> bool:
+        """Check if this node is a leaf node."""
+        return len(self.children) == 0
+
+    @property
+    def has_value(self) -> bool:
+        """Check if this node has token IDs stored."""
+        return self.token_ids is not None
+
+    def validate_token_logp_consistency(self) -> bool:
+        """Validate that token_ids, logp, and loss_mask have consistent lengths."""
+        if self.token_ids is None and self.logp is None and self.loss_mask is None:
+            return True
+
+        # Check if at least one is not None
+        if self.token_ids is not None and len(self.token_ids) > 0:
+            token_len = len(self.token_ids)
+            if self.logp is not None and len(self.logp) != token_len:
+                return False
+            if self.loss_mask is not None and len(self.loss_mask) != token_len:
+                return False
+
+        return True
+
+    @property
+    def is_evictable(self) -> bool:
+        """Check if this node can be evicted."""
+        return self.ref_count == 0 and self.token_ids is not None
+
+    def touch(self):
+        """Update access time and count."""
+        self.last_access_time = time.monotonic()
+        self.access_count += 1
+
+    def __lt__(self, other: StringTreeNode) -> bool:
+        """For heap operations - least recently used first."""
+        return self.last_access_time < other.last_access_time
+
+
+class StringRadixTrie:
+    """
+    String-based Radix Trie for efficient prefix matching and token caching.
+    Features:
+    - Efficient string prefix matching
+    - Token ID caching for matched prefixes
+    - Thread-safe operations
+    - Weight version tracking
+    - Automatic garbage collection based on weight version thresholds
+    """
+
+    def __init__(self, max_cache_size: int = 10000, gc_threshold_k: int = 5, tokenizer=None, verbose: bool = False):
+        """
+        Initialize the String Radix Trie.
+        Args:
+            max_cache_size: Maximum number of cached token IDs (triggers GC when exceeded)
+            gc_threshold_k: GC threshold - nodes with weight_version < (current_version - k) will be removed
+            tokenizer: Optional tokenizer for converting text to tokens when not found in cache
+            verbose: Whether to print debug information and tree structure
+        """
+        self.max_cache_size = max_cache_size
+        self.gc_threshold_k = gc_threshold_k
+        self.tokenizer = tokenizer
+        self.verbose = verbose
+
+        # Tree structure
+        self.root = StringTreeNode()
+        self.root.string_key = ""
+        self.root.ref_count = 1  # Root is always protected
+
+        # Cache statistics
+        self.total_entries = 0
+        self.cache_hits = 0
+        self.cache_misses = 0
+        self.cur_cache_size = 0  # Total number of token IDs across all nodes
+
+        # Thread safety
+        self._lock = threading.RLock()
+
+    def find_longest_prefix(self, text: str) -> MatchResult:
+        """
+        Find the longest cached prefix for the given text.
+        Args:
+            text: Input string to find prefix for
+        Returns:
+            MatchResult containing matched prefix, token IDs, logp, and remaining string
+        """
+        with self._lock:
+            if not text:
+                return MatchResult("", [], [], [], text, self.root)
+
+            matched_tokens = []
+            matched_logp = []
+            matched_loss_mask = []
+            matched_prefix = ""
+            current_node = self.root
+            remaining_text = text
+
+            while remaining_text:
+                # Find the best matching child that completely matches from start
+                best_child = None
+                best_key_len = 0
+
+                for child_node in current_node.children:
+                    # Only consider complete startswith matches using node's string_key
+                    if remaining_text.startswith(child_node.string_key):
+                        if len(child_node.string_key) > best_key_len:
+                            best_child = child_node
+                            best_key_len = len(child_node.string_key)
+
+                if best_child is None:
+                    # No complete startswith match found
+                    break
+
+                # Move to the best matching child
+                best_child.touch()
+                current_node = best_child
+                matched_prefix += best_child.string_key
+                remaining_text = remaining_text[best_key_len:]
+
+                # Accumulate tokens, logp, and loss_mask from this node
+                if best_child.has_value:
+                    matched_tokens.extend(best_child.token_ids)
+                    matched_logp.extend(best_child.logp)
+                    if best_child.loss_mask is not None:
+                        matched_loss_mask.extend(best_child.loss_mask)
+                    else:
+                        # If no loss_mask is stored, create default mask same as logp
+                        matched_loss_mask.extend([1] * len(best_child.token_ids))
+                    self.cache_hits += 1
+
+            if not matched_tokens:
+                self.cache_misses += 1
+
+            result = MatchResult(
+                matched_prefix, matched_tokens, matched_logp, matched_loss_mask, remaining_text, current_node
+            )
+
+            # Print tree structure if verbose is enabled
+            if self.verbose:
+                print("Tree structure after find_longest_prefix:")
+                self.pretty_print()
+
+            return result
+
+    def insert(
+        self,
+        text: str,
+        token_ids: list[int],
+        logp: list[float] | None = None,
+        loss_mask: list[int] | None = None,
+        weight_version: int | None = None,
+    ) -> bool:
+        """
+        Insert a string and its corresponding token IDs, log probabilities, and loss mask into the trie.
+        Args:
+            text: String to insert
+            token_ids: Corresponding token IDs
+            logp: Corresponding log probabilities (must match token_ids length)
+            loss_mask: Corresponding loss mask for model generation parts (must match token_ids length)
+            weight_version: Optional weight version for this insertion
+        Returns:
+            True if insertion was successful
+        """
+        with self._lock:
+            if not text or not token_ids:
+                if self.verbose:
+                    print("[RadixTree] Insertion failed: text or token_ids is empty")
+                return False
+
+            # Use provided weight version
+            current_weight_version = weight_version
+
+            # Validate logp consistency
+            if logp is not None and len(logp) != len(token_ids):
+                if self.verbose:
+                    print(
+                        f"[WARNING] Logp length {len(logp)} does not match token length {len(token_ids)} for text: {text}"
+                    )
+                    print(f"[WARNING] Logp: {logp}")
+                    print(f"[WARNING] Token IDs: {token_ids}")
+                return False
+
+            # Validate loss_mask consistency
+            if loss_mask is not None and len(loss_mask) != len(token_ids):
+                if self.verbose:
+                    print(
+                        f"[WARNING] Loss mask length {len(loss_mask)} does not match token length {len(token_ids)} for text: {text}"
+                    )
+                    print(f"[WARNING] Loss mask: {loss_mask}")
+                    print(f"[WARNING] Token IDs: {token_ids}")
+                return False
+
+            # If logp is not provided, create default values (0.0)
+            if logp is None:
+                logp = [0.0] * len(token_ids)
+
+            # If loss_mask is not provided, create default values (1 for model generation parts)
+            if loss_mask is None:
+                loss_mask = [0] * len(token_ids)
+
+            result = self._insert(text, token_ids, logp, loss_mask, current_weight_version)
+
+            # Check if GC should be triggered after insert
+            if self.cur_cache_size > self.max_cache_size and weight_version is not None:
+                if self.verbose:
+                    print(
+                        f"[RadixTree] Cache size {self.cur_cache_size} exceeds limit {self.max_cache_size}, triggering GC"
+                    )
+                gc_removed = self.gc_by_weight_version(weight_version)
+                if self.verbose:
+                    print(f"[RadixTree] GC removed {gc_removed} nodes, new cache size: {self.cur_cache_size}")
+
+            # Print tree structure if verbose is enabled
+            if self.verbose:
+                print("Tree structure after insert:")
+                self.pretty_print()
+
+            return result
+
+    def _insert(
+        self,
+        text: str,
+        token_ids: list[int],
+        logp: list[float],
+        loss_mask: list[int],
+        weight_version: int | None = None,
+    ) -> bool:
+        """Insert tokens - skip tokens for existing nodes just like we skip text."""
+
+        current_node = self.root
+        remaining_text = text
+        remaining_tokens = token_ids[:]  # Copy the tokens list
+        remaining_logp = logp[:]  # Copy the logp list
+        remaining_loss_mask = loss_mask[:]  # Copy the loss_mask list
+
+        # Track all nodes traversed during insert for weight version update
+        traversed_nodes = [current_node]
+        new_node = None
+
+        while remaining_text:
+            # Find best startswith match
+            best_child = None
+            best_key_len = 0
+
+            for child_node in current_node.children:
+                if remaining_text.startswith(child_node.string_key) and len(child_node.string_key) > best_key_len:
+                    best_child = child_node
+                    best_key_len = len(child_node.string_key)
+
+            if best_child is not None:
+                # Found existing node - skip its text and tokens
+                current_node = best_child
+                traversed_nodes.append(current_node)
+                remaining_text = remaining_text[best_key_len:]
+
+                # Skip the tokens that this existing node covers
+                if best_child.has_value:
+                    tokens_to_skip = len(best_child.token_ids)
+                    remaining_tokens = remaining_tokens[tokens_to_skip:]
+                    remaining_logp = remaining_logp[tokens_to_skip:]
+                    remaining_loss_mask = remaining_loss_mask[tokens_to_skip:]
+            else:
+                # Create new node for remaining text with remaining tokens
+                new_node = StringTreeNode()
+                new_node.parent = current_node
+                new_node.string_key = remaining_text
+
+                if remaining_tokens:  # Only assign if there are tokens left
+                    new_node.token_ids = remaining_tokens
+                    new_node.logp = remaining_logp
+                    new_node.loss_mask = remaining_loss_mask
+                    new_node.touch()
+                    # Increment cache size by number of tokens added
+                    self.cur_cache_size += len(remaining_tokens)
+
+                current_node.children.append(new_node)
+                traversed_nodes.append(new_node)
+                self.total_entries += 1
+                break
+
+        # If we've traversed the entire text and the last node doesn't have tokens,
+        # assign remaining tokens to it
+        if remaining_text == "" and not current_node.has_value:
+            if remaining_tokens:  # Only assign if there are tokens left
+                current_node.token_ids = remaining_tokens
+                current_node.logp = remaining_logp
+                current_node.loss_mask = remaining_loss_mask
+                current_node.touch()
+                self.cur_cache_size += len(remaining_tokens)
+
+        # Update weight version for all traversed nodes
+        if weight_version is not None and new_node:
+            new_node.weight_version = weight_version
+
+        return True
+
+    def remove(self, text: str) -> bool:
+        """
+        Remove a string and all nodes with this text as prefix from the trie.
+        Args:
+            text: String to remove (will also remove all strings starting with this text)
+        Returns:
+            True if any removal was performed
+        """
+        with self._lock:
+            node = self._find_node_by_text(text)
+            if node:
+                removed_count = self._clean_node_subtree(node)
+
+                # Print tree structure if verbose is enabled
+                if self.verbose:
+                    print("Tree structure after remove:")
+                    self.pretty_print()
+
+                return removed_count > 0
+            return False
+
+    def _find_node_by_text(self, text: str) -> StringTreeNode | None:
+        """
+        Find node by exact text match.
+        Args:
+            text: Text to find
+        Returns:
+            Node if found, None otherwise
+        """
+        result = self.find_longest_prefix(text)
+        if result.matched_prefix == text:
+            return result.last_node
+        return None
+
+    def _clean_node_subtree(self, node: StringTreeNode) -> int:
+        """
+        Clean a node and all its descendants.
+        This is the core cleanup function.
+        Args:
+            node: Node to clean (including all descendants)
+        Returns:
+            Number of nodes removed
+        """
+        if node == self.root:
+            return 0
+        return self._remove_node_and_descendants(node)
+
+    def _remove_node_and_descendants(self, node: StringTreeNode) -> int:
+        """
+        Remove a node and all its descendants from the trie.
+        Args:
+            node: The node to remove along with all its descendants
+        Returns:
+            Number of nodes removed
+        """
+        if node == self.root:
+            # Never remove root node
+            return 0
+
+        removed_count = 0
+
+        # First, recursively remove all descendants
+        for child in list(node.children):  # Create a copy to avoid modification during iteration
+            removed_count += self._remove_node_and_descendants(child)
+
+        # Count this node if it has data and decrement cache size
+        if node.has_value:
+            removed_count += 1
+            # Decrement cache size by number of tokens removed
+            self.cur_cache_size -= len(node.token_ids)
+
+        # Remove this node from its parent
+        if self._remove_node_from_parent(node):
+            # Update count for the node structure itself
+            pass  # _remove_node_from_parent already decrements total_entries
+
+        return removed_count
+
+    def _remove_node_from_parent(self, node: StringTreeNode) -> bool:
+        """Remove a node from its parent's children list."""
+        if node.parent and node in node.parent.children:
+            node.parent.children.remove(node)
+            self.total_entries -= 1
+            return True
+        return False
+
+    def gc_by_weight_version(self, current_weight_version: int | None = None) -> int:
+        """
+        Perform garbage collection based on weight version.
+        Remove nodes with weight_version < (current_weight_version - gc_threshold_k).
+        Args:
+            current_weight_version: Current weight version to use for GC threshold
+        Returns:
+            Number of nodes removed
+        """
+        with self._lock:
+            if current_weight_version is None:
+                if self.verbose:
+                    print("[RadixTree GC] No weight version provided, skipping GC")
+                return 0
+
+            gc_threshold = current_weight_version - self.gc_threshold_k
+            if self.verbose:
+                print(
+                    f"[RadixTree GC] Starting GC with threshold: {gc_threshold} (current_version: {current_weight_version}, k: {self.gc_threshold_k})"
+                )
+
+            nodes_to_remove = self._find_outdated_nodes(gc_threshold)
+            removed_count = 0
+
+            for node in nodes_to_remove:
+                # Validate that subtree weight versions are <= parent weight version
+                self._validate_subtree_weight_versions(node)
+                removed_count += self._clean_node_subtree(node)
+
+            if self.verbose:
+                print(f"[RadixTree GC] Completed GC, removed {removed_count} nodes")
+
+            return removed_count
+
+    def _find_outdated_nodes(self, gc_threshold: int) -> list[StringTreeNode]:
+        """
+        Find nodes that should be removed based on weight version threshold.
+        Uses layer-by-layer traversal - if parent is outdated, children are not checked.
+        Args:
+            gc_threshold: Weight version threshold (nodes < this value will be removed)
+        Returns:
+            List of nodes to remove
+        """
+        outdated_nodes = []
+
+        def check_node(node):
+            if node == self.root:
+                # Root is never removed, check its children
+                for child in node.children:
+                    check_node(child)
+                return
+
+            # Check if this node should be removed
+            if node.weight_version is not None and node.weight_version <= gc_threshold and node.has_value:
+                outdated_nodes.append(node)
+                return  # Don't check children since entire subtree will be removed
+
+            # Node is not outdated, check its children
+            for child in node.children:
+                check_node(child)
+
+        check_node(self.root)
+        return outdated_nodes
+
+    def _validate_subtree_weight_versions(self, node: StringTreeNode):
+        """
+        Validate that all nodes in subtree have weight_version <= parent weight_version.
+        Args:
+            node: Root node of subtree to validate
+        """
+
+        def validate_recursive(current_node, parent_weight_version):
+            if current_node.weight_version is not None and parent_weight_version is not None:
+                assert current_node.weight_version <= parent_weight_version, (
+                    f"Child node weight_version {current_node.weight_version} > "
+                    f"parent weight_version {parent_weight_version}"
+                )
+
+            # Recursively validate children
+            for child in current_node.children:
+                validate_recursive(child, current_node.weight_version)
+
+        # Start validation from the node itself
+        validate_recursive(node, node.weight_version)
+
+    def get_stats(self) -> dict[str, Any]:
+        """Get cache statistics."""
+        with self._lock:
+            total_requests = self.cache_hits + self.cache_misses
+            hit_rate = self.cache_hits / total_requests if total_requests > 0 else 0
+
+            return {
+                "total_entries": self.total_entries,
+                "cache_hits": self.cache_hits,
+                "cache_misses": self.cache_misses,
+                "hit_rate": hit_rate,
+                "max_cache_size": self.max_cache_size,
+                "cur_cache_size": self.cur_cache_size,
+                "gc_threshold_k": self.gc_threshold_k,
+            }
+
+    def clear(self):
+        """Clear all entries from the trie."""
+        with self._lock:
+            self.root = StringTreeNode()
+            self.root.string_key = ""
+            self.root.ref_count = 1
+            self.total_entries = 0
+            self.cache_hits = 0
+            self.cache_misses = 0
+            self.cur_cache_size = 0
+
+    def pretty_print(self):
+        """Print the trie structure in a readable format."""
+        print("String Radix Trie Structure:")
+        print("=" * 50)
+        self._print_node(self.root, 0)
+        print("=" * 50)
+        stats = self.get_stats()
+        for key, value in stats.items():
+            print(f"{key}: {value}")
+
+    def _print_node(self, node: StringTreeNode, depth: int):
+        """Recursively print node structure."""
+        indent = "  " * depth
+        key_repr = repr(node.string_key) if node.string_key else "<root>"
+        token_info = ""
+        if node.has_value:
+            token_info = f" -> tokens: {node.token_ids}"
+            if node.logp:
+                token_info += f", logp: {[round(p, 3) for p in node.logp]}"
+            if node.loss_mask:
+                token_info += f", loss_mask: {node.loss_mask}"
+        access_info = f" (accessed: {node.access_count}, ref: {node.ref_count})"
+
+        print(f"{indent}{key_repr}{token_info}{access_info}")
+
+        for child in node.children:
+            self._print_node(child, depth + 1)
+
+    def retrieve_from_text(self, text: str, return_logprob: bool = True):
+        """
+        Get tokens from text by looking up in radix tree or using tokenizer.
+        Also fetches weight version from worker during this operation.
+        Args:
+            text: Input text to get tokens for
+            return_logprob: If True, also return log probabilities
+        Returns:
+            List of token IDs corresponding to the input text if return_logprob is False.
+            Tuple of (token_ids, logp) if return_logprob is True.
+        """
+        # Call find_longest_prefix to get the match result
+        result = self.find_longest_prefix(text)
+
+        # If we have a match and it covers the entire text, return the tokens
+        if result.matched_prefix and result.token_ids:
+            additional_tokens = self.tokenizer(result.remaining_string, add_special_tokens=False)["input_ids"]
+            return (
+                result.token_ids + additional_tokens,
+                (
+                    result.logp + len(additional_tokens) * [0.0]
+                    if return_logprob
+                    else [0] * len(result.token_ids + additional_tokens)
+                ),
+                result.loss_mask + len(additional_tokens) * [0],
+            )
+        # If result is empty and input text is not empty, tokenize with tokenizer
+        # This is needed because we cannot get the prompt token id from engine response
+        # We have to manually insert the text and token into the tree
+        if self.tokenizer and text:
+            # Tokenize the text using the provided tokenizer
+            tokens = self.tokenizer(text, add_special_tokens=False)["input_ids"]
+            # Insert the text and tokens into the tree
+            self.insert(text, tokens)
+            # Return the tokens
+            return (tokens, [0.0] * len(tokens), [0] * len(tokens))
+        else:
+            raise ValueError("Tokenizer or input text can't be empty")
+
+
+# Example usage and testing
+if __name__ == "__main__":
+    # Create trie instance for testing
+    trie = StringRadixTrie(max_cache_size=100, verbose=True)
+
+    # Example usage with simplified insert
+    test_cases = [
+        ("Hello world", [1, 2, 3], [-0.1, -0.2, -0.3]),
+        ("Hello", [1, 2], [-0.1, -0.2]),
+        ("Hi there", [4, 5, 6], [-0.4, -0.5, -0.6]),
+    ]
+
+    # Insert test data with weight version and loss masks
+    print("Inserting test data...")
+    for text, tokens, logp in test_cases:
+        # Create loss_mask to match tokens length, 1 for model generation parts
+        loss_mask = [1] * len(tokens)
+        success = trie.insert(text, tokens, logp, loss_mask, weight_version=1)
+        print(f"Inserted '{text}' -> {tokens}: {success}")
+
+    print("\nTrie structure:")
+    trie.pretty_print()
+
+    # Test prefix matching
+    print("\nTesting prefix matching:")
+    test_queries = [
+        "Hello world!",  # Should match "Hello world" completely
+        "Hello everyone",  # Should match "Hello" only
+        "Hi there",  # Should match "Hi" only
+        "How are you doing?",  # Should match "How are you" completely
+        "Goodbye",  # Should not match anything
+        "Hell",  # Should not match anything (not complete startswith)
+    ]
+
+    for query in test_queries:
+        result = trie.find_longest_prefix(query)
+        print(f"Query: '{query}'")
+        print(
+            f"  Matched: '{result.matched_prefix}' -> tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}"
+        )
+        print(f"  Remaining: '{result.remaining_string}'")
+        print()
+
+    # Test removal
+    print("Testing removal:")
+    removed = trie.remove("Hello")
+    print(f"Removed 'Hello': {removed}")
+
+    result = trie.find_longest_prefix("Hello world")
+    print(
+        f"After removal - 'Hello world' -> matched: '{result.matched_prefix}', tokens: {result.token_ids}, logp: {result.logp}, loss_mask: {result.loss_mask}"
+    )
+
+    # Show final stats
+    print("\nFinal statistics:")
+    stats = trie.get_stats()
+    for key, value in stats.items():
+        print(f"{key}: {value}")
+
+    # Test GC with weight version
+    print("\nTesting GC with weight version 5:")
+    gc_removed = trie.gc_by_weight_version(5)
+    print(f"GC removed {gc_removed} nodes")
diff --git a/slime/router/middleware_hub/radix_tree_middleware.py b/slime/router/middleware_hub/radix_tree_middleware.py
new file mode 100644
index 0000000000..5df163a38d
--- /dev/null
+++ b/slime/router/middleware_hub/radix_tree_middleware.py
@@ -0,0 +1,169 @@
+import asyncio
+import json
+
+from fastapi.responses import JSONResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.requests import Request
+from starlette.responses import Response
+from transformers import AutoTokenizer
+
+from slime.utils.http_utils import post
+from slime.utils.mask_utils import get_response_lengths
+from slime.utils.types import Sample
+
+from .radix_tree import StringRadixTrie
+
+# Hop-by-hop headers that should not be forwarded
+HOP_BY_HOP = {
+    "content-length",
+    "transfer-encoding",
+    "connection",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailers",
+    "upgrade",
+}
+
+
+def _filter_headers(headers):
+    """Filter out hop-by-hop headers that should not be forwarded."""
+    return {k: v for k, v in headers.items() if k.lower() not in HOP_BY_HOP}
+
+
+async def _materialize_response(resp):
+    """Convert streaming-like Response into a regular Response/JSONResponse safely."""
+    # Collect all bytes from the streaming response
+    body = b""
+    async for chunk in resp.body_iterator:
+        body += chunk
+
+    # Try to parse as JSON based on content-type
+    ct = resp.headers.get("content-type", "")
+    headers = _filter_headers(resp.headers)
+
+    if "application/json" in ct:
+        # If it's JSON, try to parse and return as JSONResponse
+        try:
+            data = json.loads(body.decode("utf-8"))
+            return JSONResponse(content=data, status_code=resp.status_code, headers=headers)
+        except Exception:
+            # JSON parsing failed, fall back to raw bytes
+            pass
+
+    # Other types: return as raw bytes (without content-length)
+    return Response(content=body, status_code=resp.status_code, headers=headers, media_type=resp.media_type)
+
+
+class RadixTreeMiddleware(BaseHTTPMiddleware):
+    def __init__(self, app, *, router):
+        super().__init__(app)
+        self.router = router
+        self.args = router.args
+        self.tokenizer = AutoTokenizer.from_pretrained(self.args.hf_checkpoint, trust_remote_code=True)
+        self.radix_tree = StringRadixTrie(max_cache_size=10000, tokenizer=self.tokenizer, verbose=False)
+        self.router.radix_tree = self.radix_tree
+
+    async def dispatch(self, request: Request, call_next):
+
+        path = request.url.path
+
+        if path != "/generate":
+            return await call_next(request)
+
+        request_json = await request.json()
+        if "text" in request_json:
+            input_text = request_json.pop("text", "")
+        elif "input_ids" in request_json:
+            input_text = self.tokenizer.decode(request_json["input_ids"])
+        else:
+            input_text = None
+        if not input_text:
+            return await call_next(request)
+        input_tokens, input_logprobs, input_loss_mask = self.radix_tree.retrieve_from_text(
+            input_text, return_logprob=True
+        )
+        request_json["input_tokens"] = input_tokens
+        request_json["stream"] = False
+        request._json = request_json
+
+        response_data = None
+        for _ in range(5):
+            response = await call_next(request)
+
+            # If upstream returned a streaming response, materialize it to avoid Content-Length issues
+            if response.__class__.__name__ == "_StreamingResponse":
+                response = await _materialize_response(response)
+            # Try to parse JSON from the current response for meta inspection
+            try:
+                if hasattr(response, "body") and isinstance(response.body, (bytes, bytearray)):
+                    response_data = json.loads(response.body.decode("utf-8"))
+                elif hasattr(response, "content") and isinstance(response.content, (dict, list)):
+                    response_data = response.content  # JSONResponse.content is already a dict/list
+            except Exception:
+                response_data = None
+
+            if (
+                isinstance(response_data, dict)
+                and "meta_info" in response_data
+                and "finish_reason" in response_data["meta_info"]
+                and response_data["meta_info"]["finish_reason"]["type"] != "abort"
+            ):
+                break
+            # await 30 seconds for aborted responses
+            await asyncio.sleep(30)
+
+        if isinstance(response_data, dict) and "text" in response_data and "output_ids" in response_data:
+            generated_text = response_data["text"]
+
+            full_text = input_text + generated_text
+            if full_text:
+                try:
+                    if "output_token_logprobs" in response_data.get("meta_info", {}):
+                        generated_token_logprobs = [
+                            item[0] for item in response_data["meta_info"]["output_token_logprobs"]
+                        ]
+                        generated_token_ids = [item[1] for item in response_data["meta_info"]["output_token_logprobs"]]
+                        full_logprobs = input_logprobs + generated_token_logprobs
+                        full_token_ids = input_tokens + generated_token_ids
+                        full_loss_mask = input_loss_mask + [1] * len(generated_token_ids)
+                        self.radix_tree.insert(
+                            full_text,
+                            full_token_ids,
+                            full_logprobs,
+                            full_loss_mask,
+                            weight_version=response_data["meta_info"]["weight_version"],
+                        )
+                    else:
+                        generated_token_ids = self.tokenizer(generated_text, add_special_tokens=False)["input_ids"]
+                        full_token_ids = input_tokens + generated_token_ids
+                        full_loss_mask = input_loss_mask + [1] * len(generated_token_ids)
+                        self.radix_tree.insert(
+                            full_text,
+                            full_token_ids,
+                            None,
+                            full_loss_mask,
+                            weight_version=response_data["meta_info"]["weight_version"],
+                        )
+
+                    if getattr(self.router, "verbose", False):
+                        print(f"[slime-router] Successfully cached trajectory with {len(full_token_ids)} tokens")
+                except Exception as e:
+                    if getattr(self.router, "verbose", False):
+                        print(f"[slime-router] Warning: Failed to cache trajectory: {e}")
+        return response
+
+
+async def postprocess_sample_with_radix_tree(args, sample: Sample, output: dict):
+    assert not args.partial_rollout, "Currently partial rollout is not supported when using slime router"
+    retrieve_url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/retrieve_from_text"
+    retrieve_payload = {"text": sample.prompt + output["text"], "return_logp": True}
+    retrieve_output = await post(retrieve_url, retrieve_payload)
+    sample.tokens = retrieve_output["tokens"]
+    sample.response += output["text"]
+    sample.loss_mask = retrieve_output["loss_mask"]
+    sample.response_length = get_response_lengths([sample.loss_mask])[0]
+    sample.loss_mask = sample.loss_mask[-sample.response_length :]
+    sample.rollout_log_probs = retrieve_output["rollout_logp"][-sample.response_length :]
+    return sample
diff --git a/slime/router/router.py b/slime/router/router.py
index b77dd975f4..669094e442 100644
--- a/slime/router/router.py
+++ b/slime/router/router.py
@@ -45,7 +45,7 @@ def __init__(self, args, verbose=False):
         max_connections = getattr(args, "slime_router_max_connections", None)
         if max_connections is None:
             max_connections = (
-                args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
+                args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
             )
 
         timeout = getattr(args, "slime_router_timeout", None)
diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py
index f9617e1f5e..fb2d305a63 100644
--- a/slime/utils/arguments.py
+++ b/slime/utils/arguments.py
@@ -1,12 +1,14 @@
 import argparse
-import copy
 import json
 import logging
 import os
 from typing import Any
 
 import yaml
+from sglang_router.launch_router import RouterArgs
 
+from slime.backends.sglang_utils.arguments import sglang_parse_args
+from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args
 from slime.utils.eval_config import EvalDatasetConfig, build_eval_dataset_configs, ensure_dataset_list
 from slime.utils.logging_utils import configure_logger
 
@@ -37,6 +39,12 @@ def add_cluster_arguments(parser):
             parser.add_argument(
                 "--actor-num-gpus-per-node", type=int, default=8, help="Number of gpus per node for training actor"
             )
+            parser.add_argument(
+                "--critic-num-nodes", type=int, default=None, help="Number of nodes for training actor"
+            )
+            parser.add_argument(
+                "--critic-num-gpus-per-node", type=int, default=None, help="Number of gpus per node for training actor"
+            )
 
             parser.add_argument(
                 "--rollout-num-gpus",
@@ -109,6 +117,12 @@ def add_train_arguments(parser):
                 default="thd",
                 help="The qkv layout for Megatron backend.",
             )
+            parser.add_argument(
+                "--true-on-policy-mode",
+                action="store_true",
+                default=False,
+                help="Whether to enable true-on-policy mode.",
+            )
             parser.add_argument(
                 "--train-env-vars",
                 type=json.loads,
@@ -224,45 +238,6 @@ def add_rollout_arguments(parser):
                     "Also, sometimes this will help alleviate the bug that transformers cannot find certain model."
                 ),
             )
-            parser.add_argument(
-                "--rollout-backend",
-                type=str,
-                choices=["sglang", "vllm"],
-                default="sglang",
-                help="Rollout inference backend: sglang (default) or vllm.",
-            )
-            parser.add_argument(
-                "--vllm-base-url",
-                type=str,
-                default=None,
-                help="vLLM server URL (set automatically when using managed vLLM).",
-            )
-            parser.add_argument(
-                "--vllm-max-retries",
-                type=int,
-                default=3,
-                help="Max HTTP retries for vLLM requests.",
-            )
-            parser.add_argument(
-                "--vllm-weight-sync-packed",
-                action=argparse.BooleanOptionalAction,
-                default=True,
-                help="Use vLLM packed weight transfer for non-colocate (default: True). Disable for per-bucket mode.",
-            )
-            parser.add_argument(
-                "--vllm-gpu-memory-utilization",
-                type=float,
-                default=0.4,
-                help="Fraction of GPU memory for vLLM KV cache (default: 0.85). "
-                     "Lower this if the training model leaves insufficient free memory.",
-            )
-            parser.add_argument(
-                "--rollout-server-concurrency",
-                type=int,
-                default=512,
-                help="Maximum number of concurrent requests sent to the rollout server. "
-                     "Controls request parallelism for any backend (sglang or vllm).",
-            )
             parser.add_argument(
                 "--rollout-function-path",
                 type=str,
@@ -781,21 +756,16 @@ def add_algo_arguments(parser):
             reset_arg(parser, "--calculate-per-token-loss", action="store_true")
             reset_arg(parser, "--lr", type=float, default=1e-6)
 
+            parser.add_argument("--num-critic-only-steps", type=int, default=0, help="Number of critic only steps")
+            parser.add_argument("--critic-load", type=str, default=None, help="The checkpoint for critic model.")
+            parser.add_argument("--critic-save", type=str, default=None, help="The checkpoint for critic model.")
+            parser.add_argument("--critic-lr", type=float, default=None, help="The lr for critic model")
+            parser.add_argument("--critic-train-only", action="store_true", default=False, help="Only train critic")
             parser.add_argument(
-                "--num-critic-only-steps",
+                "--critic-lr-warmup-iters",
                 type=int,
                 default=0,
-                help="Number of initial rollout steps that train critic only; set >= num_rollout for critic-only runs",
-            )
-            parser.add_argument(
-                "--megatron-config-path",
-                type=str,
-                default=None,
-                help=(
-                    "Path to a structured YAML config for Megatron roles. The file should use "
-                    "a top-level 'megatron' key with role-tagged entries; the critic runtime will "
-                    "select exactly one entry with role=critic. Legacy 'critic' configs are still accepted."
-                ),
+                help="number of iterations to linearly warmup for critic model.",
             )
 
             parser.add_argument("--eps-clip", type=float, default=0.2, help="PPO clip range")
@@ -865,19 +835,6 @@ def add_algo_arguments(parser):
                     "This is useful for sft or custom loss function."
                 ),
             )
-            parser.add_argument(
-                "--custom-advantage-function-path",
-                type=str,
-                default=None,
-                help=(
-                    "Path to a custom advantage/returns computation function. "
-                    "When set, this function replaces the built-in compute_advantages_and_returns. "
-                    "Signature: def custom_fn(args, rollout_data) -> None. "
-                    "The function should set rollout_data['advantages'] and rollout_data['returns'] in-place. "
-                    "Critic values are available in rollout_data['values']. "
-                    "(e.g., my_module.py:my_advantage_fn)."
-                ),
-            )
             parser.add_argument(
                 "--use-kl-loss", action="store_true", default=False, help="whether to use KL loss from GRPO"
             )
@@ -1080,13 +1037,7 @@ def add_router_arguments(parser):
                 default=3,
                 help="Number of consecutive failures before marking a worker as unhealthy.",
             )
-            # try:
-            #     from sglang_router.launch_router import RouterArgs
-
-            #     RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
-            # except ImportError:
-            #     pass  # sglang not installed — router args unavailable (vllm-only mode)
-            # RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
+            RouterArgs.add_cli_args(parser, use_router_prefix=True, exclude_host_port=True)
             return parser
 
         # wandb
@@ -1149,7 +1100,7 @@ def add_wandb_arguments(parser):
                 default=None,
                 help=(
                     "Log statistics of the category of reward, such as why the reward function considers it as failed. "
-                    "Specify the key in the reward dict using this argument."
+                    "Specify the key in the reward dict using this argument.",
                 ),
             )
             parser.add_argument(
@@ -1331,7 +1282,7 @@ def add_rollout_buffer_arguments(parser):
                 "--loss-mask-type",
                 type=str,
                 default="qwen",
-                choices=["qwen", "qwen3", "qwen3_5", "distill_qwen"],
+                choices=["qwen", "qwen3", "distill_qwen"],
                 help="Loss mask type",
             )
             parser.add_argument(
@@ -1476,11 +1427,10 @@ def _pre_parse_mode():
     the final ``args`` after Phase 2 parsing.
     """
     temp_parser = argparse.ArgumentParser(add_help=False, allow_abbrev=False)
-    temp_parser.add_argument("--train-backend", type=str, choices=["megatron"], default="megatron")
+    temp_parser.add_argument("--train-backend", type=str, choices=["megatron", "fsdp"], default="megatron")
     temp_parser.add_argument("--debug-rollout-only", action="store_true", default=False)
     temp_parser.add_argument("--debug-train-only", action="store_true", default=False)
     temp_parser.add_argument("--load-debug-rollout-data", type=str, default=None)
-    temp_parser.add_argument("--rollout-backend", type=str, choices=["sglang", "vllm"], default="sglang")
     temp_args, _ = temp_parser.parse_known_args()
     return temp_args
 
@@ -1492,30 +1442,33 @@ def parse_args(add_custom_arguments=None):
     add_slime_arguments = get_slime_extra_args_provider(add_custom_arguments)
 
     pre = _pre_parse_mode()
-    skip_sglang = (
-        pre.debug_train_only
-        or pre.load_debug_rollout_data is not None
-        or pre.rollout_backend == "vllm"
-    )
+    skip_sglang = pre.debug_train_only or pre.load_debug_rollout_data is not None
 
     # Phase 1: Parse sglang args independently (separate parser, parse_known_args).
-    # Skipped when sglang servers are not needed or when using vLLM backend.
+    # Skipped when sglang servers are not needed.
     sglang_ns = None
     if not skip_sglang:
-        from slime.backends.sglang_utils.arguments import sglang_parse_args
-
         sglang_ns = sglang_parse_args()
 
-    # Phase 2: Parse megatron + slime args.
+    # Phase 2: Parse megatron/fsdp + slime args.
     # Uses ignore_unknown_args=True so that --sglang-* and pre-parsed CLI flags
-    # are silently ignored by the megatron parser.
-    from slime.backends.megatron_utils.arguments import megatron_parse_args
-    from slime.backends.megatron_utils.arguments import validate_args as megatron_validate_args
+    # are silently ignored by the megatron/fsdp parser.
+    if pre.train_backend == "megatron":
+        from slime.backends.megatron_utils.arguments import megatron_parse_args
+        from slime.backends.megatron_utils.arguments import validate_args as megatron_validate_args
+
+        args = megatron_parse_args(
+            extra_args_provider=add_slime_arguments,
+            skip_hf_validate=pre.debug_rollout_only,
+        )
+    else:
+        logger.warning(
+            "🚧 🚧 🚧 FSDP backend is being rewritten, please use Megatron backend for better stability. 🚧 🚧 🚧"
+        )
 
-    args = megatron_parse_args(
-        extra_args_provider=add_slime_arguments,
-        skip_hf_validate=pre.debug_rollout_only,
-    )
+        from slime.backends.fsdp_utils.arguments import fsdp_parse_args
+
+        args = fsdp_parse_args(extra_args_provider=add_slime_arguments, ignore_unknown_args=True)
 
     # Merge pre-parsed args into the main namespace
     for key, value in vars(pre).items():
@@ -1531,103 +1484,12 @@ def parse_args(add_custom_arguments=None):
     if pre.train_backend == "megatron" and not args.debug_rollout_only:
         megatron_validate_args(args)
 
-    if not args.debug_train_only and getattr(args, "rollout_backend", "sglang") == "sglang":
-        from slime.backends.sglang_utils.arguments import validate_args as sglang_validate_args
-
+    if not args.debug_train_only:
         sglang_validate_args(args)
-    elif getattr(args, "rollout_backend", "sglang") == "vllm":
-        # Set sglang aliases that the rest of the codebase expects
-        args.sglang_dp_size = getattr(args, "sglang_data_parallel_size", 1) or 1
-        args.sglang_pp_size = getattr(args, "sglang_pipeline_parallel_size", 1) or 1
-        args.sglang_ep_size = getattr(args, "sglang_expert_parallel_size", 1) or 1
-        args.sglang_tp_size = args.rollout_num_gpus_per_engine
-        if not hasattr(args, "sglang_speculative_algorithm"):
-            args.sglang_speculative_algorithm = None
 
     return args
 
 
-def _apply_megatron_role_overrides(base_args, overrides, role):
-    role_args = copy.deepcopy(base_args)
-    ignored_keys = {"num_nodes", "num_gpus_per_node"}
-
-    # Apply overrides from the YAML config.
-    # Unspecified keys inherit from base_args via deepcopy.
-    for key, value in overrides.items():
-        if key in ignored_keys:
-            logger.info(f"Ignoring {role} config key '{key}'; GPU allocation always follows CLI args.")
-            continue
-        if not hasattr(role_args, key):
-            logger.warning(f"{role.capitalize()} config key '{key}' is not a known argument, setting it anyway.")
-        else:
-            # YAML safe_load doesn't parse scientific notation (e.g. 1e-5) as float.
-            # Coerce the value to match the type of the existing attribute.
-            original = getattr(role_args, key)
-            if original is not None and isinstance(value, str) and isinstance(original, (int, float)):
-                try:
-                    value = type(original)(value)
-                except (ValueError, TypeError):
-                    pass
-        setattr(role_args, key, value)
-
-    if role == "critic":
-        # Critic-specific: disable features that only apply to actors.
-        role_args.kl_coef = 0
-        role_args.use_opd = False
-        role_args.custom_advantage_function_path = None
-        role_args.untie_embeddings_and_output_weights = True
-
-    return role_args
-
-
-def parse_megatron_role_args(base_args, megatron_config_path, role):
-    """Parse role-specific arguments from a unified Megatron YAML config.
-
-    The config must contain a top-level ``megatron`` list with per-role entries.
-    Missing roles inherit the base args unchanged.
-    """
-    assert role in {"actor", "critic"}, f"Unsupported Megatron config role: {role}"
-
-    with open(megatron_config_path) as f:
-        raw_config = yaml.safe_load(f) or {}
-
-    assert "megatron" in raw_config, (
-        "megatron config must contain a top-level 'megatron' list, e.g. "
-        "megatron: [{name: default, role: actor, overrides: {...}}]"
-    )
-
-    overrides = {}
-    megatron_entries = raw_config["megatron"]
-    assert isinstance(megatron_entries, list), (
-        "megatron config 'megatron' field must be a list, e.g. "
-        "megatron: [{name: default, role: actor, overrides: {...}}]"
-    )
-    role_entries = [entry for entry in megatron_entries if entry.get("role") == role]
-    assert len(role_entries) <= 1, (
-        f"megatron config must contain at most one entry with role={role}, e.g. "
-        f"megatron: [{{name: default, role: {role}, overrides: {{...}}}}]"
-    )
-    if role_entries:
-        role_entry = role_entries[0]
-        overrides = role_entry.get("overrides") or role_entry.get("args") or {}
-    else:
-        logger.info(
-            f"No megatron config entry with role={role} found in {megatron_config_path}; using inherited args."
-        )
-
-    role_args = _apply_megatron_role_overrides(base_args, overrides, role)
-    logger.info(
-        f"Parsed megatron config for role={role} from {megatron_config_path}: overrides = {list(overrides.keys())}"
-    )
-
-    return role_args
-
-
-def parse_critic_args(actor_args, megatron_config_path):
-    """Backward-compatible wrapper for critic-specific Megatron role parsing."""
-    return parse_megatron_role_args(actor_args, megatron_config_path, role="critic")
-
-
 def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
     """
     Build evaluation dataset configurations from either --eval-config or --eval-prompt-data.
@@ -1674,13 +1536,6 @@ def _resolve_eval_datasets(args) -> list[EvalDatasetConfig]:
 def slime_validate_args(args):
     args.eval_datasets = _resolve_eval_datasets(args)
 
-    if args.use_slime_router:
-        logger.warning(
-            "--use-slime-router is deprecated and ignored. slime now always uses sglang_router "
-            "built from https://github.com/zhuzilin/sgl-router."
-        )
-        args.use_slime_router = False
-
     if args.kl_coef != 0 or args.use_kl_loss:
         if not os.path.exists(args.ref_load):
             raise FileNotFoundError(f"ref_load {args.ref_load} does not exist, please check the path.")
@@ -1800,9 +1655,14 @@ def slime_validate_args(args):
         args.debug_train_only = True
 
     args.use_critic = args.advantage_estimator == "ppo"
-    # Critic always uses the same GPU count as actor.
-    args.critic_num_gpus_per_node = args.actor_num_gpus_per_node
-    args.critic_num_nodes = args.actor_num_nodes
+    if args.critic_num_gpus_per_node is None:
+        args.critic_num_gpus_per_node = args.actor_num_gpus_per_node
+    if args.critic_num_nodes is None:
+        args.critic_num_nodes = args.actor_num_nodes
+    if args.critic_load is None:
+        args.critic_load = args.load
+    if args.critic_lr is None:
+        args.critic_lr = args.lr
 
     if args.offload:
         args.offload_train = True
@@ -1844,10 +1704,6 @@ def slime_validate_args(args):
         args.offload_train = False
     if args.offload_rollout is None:
         args.offload_rollout = False
-    args.offload_rollout = False 
-
-    if args.use_critic:
-        args.offload_train = True
 
     if args.eval_function_path is None:
         args.eval_function_path = args.rollout_function_path
diff --git a/slime/utils/data.py b/slime/utils/data.py
index 4bb81e5677..c3dbfa96a0 100644
--- a/slime/utils/data.py
+++ b/slime/utils/data.py
@@ -89,30 +89,15 @@ def filter_long_prompt(origin_samples: list[Sample], tokenizer, processor, max_l
         return origin_samples
 
     if processor:
-        # Use processor only for samples with actual multimodal content; use batched tokenizer for text-only.
-        text_only = []
-        multimodal = []
-        for sample in origin_samples:
-            if sample.multimodal_inputs and any(v is not None for v in sample.multimodal_inputs.values()):
-                multimodal.append(sample)
-            else:
-                text_only.append(sample)
         filtered_samples = []
-        if text_only:
-            prompts = [s.prompt for s in text_only]
-            input_ids_list = tokenizer(prompts, add_special_tokens=False)["input_ids"]
-            for sample, input_ids in zip(text_only, input_ids_list, strict=True):
-                if len(input_ids) <= max_length:
-                    filtered_samples.append(sample)
-        if multimodal:
+        for sample in origin_samples:
             from slime.utils.processing_utils import process_vision_info
 
-            for sample in multimodal:
-                multimodal_inputs = process_vision_info(sample.prompt, processor)
-                processor_output = processor(text=sample.prompt, **multimodal_inputs)
-                input_ids = processor_output["input_ids"][0]
-                if len(input_ids) <= max_length:
-                    filtered_samples.append(sample)
+            multimodal_inputs = process_vision_info(sample.prompt, processor)
+            processor_output = processor(text=sample.prompt, **multimodal_inputs)
+            input_ids = processor_output["input_ids"][0]
+            if len(input_ids) <= max_length:
+                filtered_samples.append(sample)
     else:
         prompts = [sample.prompt for sample in origin_samples]
         input_ids_list = tokenizer(prompts, add_special_tokens=False)["input_ids"]
diff --git a/slime/utils/distributed_utils.py b/slime/utils/distributed_utils.py
index 7f7e57271a..7918657a99 100644
--- a/slime/utils/distributed_utils.py
+++ b/slime/utils/distributed_utils.py
@@ -3,7 +3,6 @@
 
 import torch
 import torch.distributed as dist
-from packaging.version import parse as V
 from torch.distributed.distributed_c10d import (
     Backend,
     PrefixStore,
@@ -14,6 +13,7 @@
     rendezvous,
 )
 
+
 GLOO_GROUP = None
 
 
@@ -74,7 +74,7 @@ def init_process_group(
     # NOTE: The pg_options parameter was renamed into backend_options in PyTorch 2.6.0
     # https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844
     # We need to determine the appropriate parameter name based on PyTorch version
-    pg_options_param_name = "backend_options" if V(torch.__version__) >= V("2.6") else "pg_options"
+    pg_options_param_name = "backend_options" if str(torch.__version__) >= "2.6" else "pg_options"
     pg, _ = _new_process_group_helper(
         world_size,
         rank,
diff --git a/slime/utils/eval_config.py b/slime/utils/eval_config.py
index 8342c7b20e..c3c20db1b8 100644
--- a/slime/utils/eval_config.py
+++ b/slime/utils/eval_config.py
@@ -118,12 +118,6 @@ class EvalDatasetConfig:
     # per-dataset custom generate function (e.g., for tool calling)
     custom_generate_function_path: str | None = None
 
-    # app_service URL for server mode generation (e.g., "http://localhost:18080")
-    # If set, eval will use ServerGenerationProxy to generate through AppServer
-    app_service: str | None = None
-
-    eval_task_timeout: int | None = None
-
     metadata_overrides: dict[str, Any] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
diff --git a/slime/utils/external_utils/command_utils.py b/slime/utils/external_utils/command_utils.py
index cdf2cbe0b0..9f51ecdf20 100644
--- a/slime/utils/external_utils/command_utils.py
+++ b/slime/utils/external_utils/command_utils.py
@@ -28,7 +28,6 @@ def convert_checkpoint(
     hf_checkpoint: str | None = None,
 ):
     hf_checkpoint = hf_checkpoint or f"/root/models/{model_name}"
-    normalized_extra_args = f" {extra_args.strip()}" if extra_args.strip() else ""
 
     # TODO shall we make it in host-mapped folder and thus can cache it to speedup CI
     path_dst = f"{dir_dst}/{model_name}_torch_dist"
@@ -60,7 +59,7 @@ def convert_checkpoint(
         "${MODEL_ARGS[@]} "
         f"--hf-checkpoint {hf_checkpoint} "
         f"--save {path_dst}"
-        f"{normalized_extra_args}"
+        f"{extra_args}"
     )
 
 
@@ -107,6 +106,9 @@ def execute_train(
     external_ray = get_bool_env_var("SLIME_SCRIPT_EXTERNAL_RAY")
     master_addr = os.environ.get("MASTER_ADDR", "127.0.0.1")
 
+    train_backend_fsdp = "--train-backend fsdp" in train_args
+    assert train_backend_fsdp == (megatron_model_type is None)
+
     exec_command(
         "pkill -9 sglang; "
         "sleep 3; "
@@ -138,7 +140,14 @@ def execute_train(
         {
             "env_vars": {
                 "PYTHONPATH": "/root/Megatron-LM/",
-                "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+                # If setting this in FSDP, the computation communication overlapping may have issues
+                **(
+                    {}
+                    if train_backend_fsdp
+                    else {
+                        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+                    }
+                ),
                 "NCCL_NVLS_ENABLE": str(int(check_has_nvlink())),
                 "no_proxy": f"127.0.0.1,{master_addr}",
                 # This is needed by megatron / torch distributed in multi-node setup
diff --git a/slime/utils/health_monitor.py b/slime/utils/health_monitor.py
index 52190267af..975330963f 100644
--- a/slime/utils/health_monitor.py
+++ b/slime/utils/health_monitor.py
@@ -20,8 +20,8 @@ class RolloutHealthMonitor:
     - stop(): Stop the monitor thread completely (called during dispose)
     """
 
-    def __init__(self, server_group, args):
-        self._server_group = server_group
+    def __init__(self, engine_group, args):
+        self._engine_group = engine_group
 
         self._thread = None
         self._stop_event = None
@@ -38,7 +38,7 @@ def start(self) -> bool:
         Returns:
             True if the monitor was started, False if there are no engines to monitor.
         """
-        if not self._server_group.all_engines:
+        if not self._engine_group.all_engines:
             return False
 
         if self._thread is not None:
@@ -135,7 +135,7 @@ def _health_monitor_loop(self) -> None:
                 break
 
     def _run_health_checks(self) -> None:
-        for rollout_engine_id, engine in enumerate(self._server_group.engines):
+        for rollout_engine_id, engine in enumerate(self._engine_group.engines):
             if self._stop_event is not None and self._stop_event.is_set():
                 break
             if self._pause_event is not None and self._pause_event.is_set():
@@ -158,12 +158,12 @@ def _check_engine_health(self, rollout_engine_id, engine) -> None:
             logger.debug(f"Health check passed for rollout engine {rollout_engine_id}")
 
     def _kill_engine(self, rollout_engine_id: int):
-        logger.info(f"Killing server group {rollout_engine_id}...")
+        logger.info(f"Killing engine group {rollout_engine_id}...")
         for i in range(
-            rollout_engine_id * self._server_group.nodes_per_engine,
-            (rollout_engine_id + 1) * self._server_group.nodes_per_engine,
+            rollout_engine_id * self._engine_group.nodes_per_engine,
+            (rollout_engine_id + 1) * self._engine_group.nodes_per_engine,
         ):
-            engine = self._server_group.all_engines[i]
+            engine = self._engine_group.all_engines[i]
             if engine:
                 logger.info(f"Shutting down and killing engine at index {i}")
                 try:
@@ -174,4 +174,4 @@ def _kill_engine(self, rollout_engine_id: int):
                     logger.warning(f"Fail to kill engine at index {i} (e: {e})")
             else:
                 logger.info(f"Engine at index {i} is already None")
-            self._server_group.all_engines[i] = None
+            self._engine_group.all_engines[i] = None
diff --git a/slime/utils/http_utils.py b/slime/utils/http_utils.py
index d68e340796..d7807f3b7a 100644
--- a/slime/utils/http_utils.py
+++ b/slime/utils/http_utils.py
@@ -204,12 +204,11 @@ def init_http_client(args):
     if not args.rollout_num_gpus:
         return
 
-    _client_concurrency = args.rollout_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
+    _client_concurrency = args.sglang_server_concurrency * args.rollout_num_gpus // args.rollout_num_gpus_per_engine
     if _http_client is None:
         _http_client = httpx.AsyncClient(
             limits=httpx.Limits(max_connections=_client_concurrency),
             timeout=httpx.Timeout(None),
-            trust_env=False,  # internal SGLang comm only — never route through system proxy
         )
 
     # Optionally initialize distributed POST via Ray without changing interfaces
@@ -221,8 +220,8 @@ def init_http_client(args):
 def _init_ray_distributed_post(args):
     """Initialize one or more Ray async actors per node for HTTP POST.
 
-    Uses NodeAffinitySchedulingStrategy to place actors on alive Ray nodes.
-    Creates ``args.num_gpus_per_node`` actors per node.
+    Uses NodeAffinitySchedulingStrategy to place actors on distinct nodes.
+    Controlled by SLIME_HTTP_POST_ACTORS_PER_NODE.
     """
     global _post_actors
     if _post_actors:
@@ -244,7 +243,6 @@ def __init__(self, concurrency: int):
             self._client = httpx.AsyncClient(
                 limits=httpx.Limits(max_connections=max(1, concurrency)),
                 timeout=httpx.Timeout(None),
-                trust_env=False,  # internal SGLang comm only — never route through system proxy
             )
 
         async def do_post(self, url, payload, max_retries=60, headers=None):
@@ -252,8 +250,8 @@ async def do_post(self, url, payload, max_retries=60, headers=None):
 
     # Create actors per node
     created = []
-    total_actors = max(1, len(nodes) * args.num_gpus_per_node)
-    per_actor_conc = max(1, (_client_concurrency + total_actors - 1) // total_actors)
+    # Distribute client concurrency across actors (at least 1 per actor)
+    per_actor_conc = (_client_concurrency + len(nodes)) // len(nodes)
 
     for node in nodes:
         node_id = node["NodeID"]
@@ -276,16 +274,13 @@ async def post(url, payload, max_retries=60, headers=None):
     # If distributed mode is enabled and actors exist, dispatch via Ray.
     if _distributed_post_enabled and _post_actors:
         try:
+            import ray
+
             actor = _next_actor()
             if actor is not None:
-                # Await the Ray ObjectRef directly. The previous
-                # `asyncio.to_thread(ray.get, obj_ref)` blocked an OS thread
-                # from the default ThreadPoolExecutor (capped at
-                # `min(32, cpu+4)`), which becomes a hard upper bound on the
-                # number of in-flight POSTs that can be waited on in parallel
-                # and produces large tail latencies under high concurrency.
+                # Use a thread to avoid blocking the event loop on ray.get
                 obj_ref = actor.do_post.remote(url, payload, max_retries, headers=headers)
-                return await obj_ref
+                return await asyncio.to_thread(ray.get, obj_ref)
         except Exception as e:
             logger.info(f"[http_utils] Distributed POST failed, falling back to local: {e} (url={url})")
             # fall through to local
diff --git a/slime/utils/logging_utils.py b/slime/utils/logging_utils.py
index 11348a4074..5fc0fad357 100644
--- a/slime/utils/logging_utils.py
+++ b/slime/utils/logging_utils.py
@@ -31,20 +31,6 @@ def init_tracking(args, primary: bool = True, **kwargs):
         wandb_utils.init_wandb_secondary(args, **kwargs)
 
 
-def update_tracking_open_metrics(args, router_addr):
-    wandb_utils.reinit_wandb_primary_with_open_metrics(args, router_addr)
-
-
-def finish_tracking(args):
-    if not args.use_wandb:
-        return
-    try:
-        if wandb.run is not None:
-            wandb.finish()
-    except Exception:
-        logging.getLogger(__name__).exception("Failed to finish wandb run")
-
-
 # TODO further refactor, e.g. put TensorBoard init to the "init" part
 def log(args, metrics, step_key: str):
     if args.use_wandb:
diff --git a/slime/utils/mask_utils.py b/slime/utils/mask_utils.py
index efe5e159f1..0ddb3a1410 100644
--- a/slime/utils/mask_utils.py
+++ b/slime/utils/mask_utils.py
@@ -9,11 +9,8 @@ def get_response_lengths(loss_masks: list[list[int]]) -> list[int]:
 class MultiTurnLossMaskGenerator:
     def __init__(self, tokenizer: AutoTokenizer, tokenizer_type: str = "qwen"):
         self.tokenizer = tokenizer
+        self.system_message_length, self.gen_token_length = self.get_system_message_length()
         self.tokenizer_type = tokenizer_type
-        self.system_message_length = 0
-        self.gen_token_length = 0
-        if self.tokenizer_type in ("qwen", "qwen3"):
-            self.system_message_length, self.gen_token_length = self.get_system_message_length()
 
     def get_response_lengths(self, loss_masks: list[list[int]]) -> list[int]:
         return get_response_lengths(loss_masks)
@@ -41,11 +38,7 @@ def get_system_message_length(self) -> tuple[int, int]:
         end_interval = len(chat_template_token_ids) - len(raw_token_ids) - idx_2
         gen_token_length = len(
             self.tokenizer.apply_chat_template(
-                test_messages,
-                add_special_tokens=False,
-                tokenize=True,
-                add_generation_prompt=True,
-                return_dict=False,
+                test_messages, add_special_tokens=False, tokenize=True, add_generation_prompt=True
             )
         ) - len(chat_template_token_ids)
 
@@ -60,11 +53,9 @@ def gen_multi_turn_loss_mask_qwen(
 
         for i, message in enumerate(messages):
             if i == 0:
-                message_ids = self.tokenizer.apply_chat_template(
-                    [message], tokenize=True, tools=tools, return_dict=False
-                )
+                message_ids = self.tokenizer.apply_chat_template([message], tokenize=True, tools=tools)
             else:
-                message_ids = self.tokenizer.apply_chat_template([message], tokenize=True, return_dict=False)
+                message_ids = self.tokenizer.apply_chat_template([message], tokenize=True)
 
             if message["role"] != "system" and i > 0:
                 message_ids = message_ids[self.system_message_length :]
@@ -89,23 +80,16 @@ def gen_multi_turn_loss_mask_qwen3(
         all_token_ids = []
 
         prefix_message = {"role": "user", "content": "FOR CALCULATING LOSS MASK ONLY"}
-        prefix_token_ids = self.tokenizer.apply_chat_template([prefix_message], tokenize=True, return_dict=False)
+        prefix_token_ids = self.tokenizer.apply_chat_template([prefix_message], tokenize=True)
 
         for i, message in enumerate(messages):
             if i == 0:
                 tailed_message_ids = self.tokenizer.apply_chat_template(
-                    [message, prefix_message],
-                    tokenize=True,
-                    tools=tools,
-                    return_dict=False,
+                    [message, prefix_message], tokenize=True, tools=tools
                 )
                 message_ids = tailed_message_ids[: -len(prefix_token_ids)]
             else:
-                prefixed_message_ids = self.tokenizer.apply_chat_template(
-                    [prefix_message, message],
-                    tokenize=True,
-                    return_dict=False,
-                )
+                prefixed_message_ids = self.tokenizer.apply_chat_template([prefix_message, message], tokenize=True)
                 message_ids = prefixed_message_ids[len(prefix_token_ids) :]
 
             if message["role"] != "system" and i > 0:
@@ -124,77 +108,6 @@ def gen_multi_turn_loss_mask_qwen3(
 
         return all_token_ids, all_loss_masks
 
-    def gen_multi_turn_loss_mask_qwen3_5(
-        self, messages: list[dict], tools: list[dict] = None
-    ) -> tuple[list[int], list[int]]:
-        rendered_text = self.tokenizer.apply_chat_template(messages, tokenize=False, tools=tools, return_dict=False)
-        tokenized = self.tokenizer(rendered_text, add_special_tokens=False, return_offsets_mapping=True)
-        token_ids = tokenized["input_ids"]
-        offset_mapping = tokenized.get("offset_mapping")
-
-        if offset_mapping is None:
-            raise ValueError(
-                "Qwen3.5 loss mask generation requires a fast tokenizer " "with `return_offsets_mapping` support."
-            )
-
-        expected_token_ids = self.tokenizer.apply_chat_template(
-            messages, tokenize=True, tools=tools, return_dict=False
-        )
-        if token_ids != expected_token_ids:
-            raise ValueError(
-                "Qwen3.5 rendered text tokenization does not match "
-                "`apply_chat_template(..., tokenize=True)` output."
-            )
-
-        assistant_header = "<|im_start|>assistant\n"
-        think_prefix = "<think>\n"
-        end_marker = "<|im_end|>"
-
-        char_mask = [0] * len(rendered_text)
-        cursor = 0
-
-        for message in messages:
-            if message["role"] != "assistant":
-                continue
-
-            header_pos = rendered_text.find(assistant_header, cursor)
-            if header_pos < 0:
-                raise ValueError("Failed to locate assistant message in rendered Qwen3.5 chat template output.")
-
-            content_start = header_pos + len(assistant_header)
-            end_pos = rendered_text.find(end_marker, content_start)
-            if end_pos < 0:
-                raise ValueError("Failed to locate <|im_end|> for assistant message in rendered text.")
-
-            span_end = end_pos + len(end_marker)
-            if span_end < len(rendered_text) and rendered_text[span_end] == "\n":
-                span_end += 1
-            cursor = span_end
-
-            if message.get("step_loss_mask", 1) != 1:
-                continue
-
-            if rendered_text[content_start : content_start + len(think_prefix)] == think_prefix:
-                mask_start = content_start + len(think_prefix)
-            else:
-                mask_start = content_start
-
-            for pos in range(mask_start, span_end):
-                char_mask[pos] = 1
-
-        char_mask_prefix_sum = [0]
-        for value in char_mask:
-            char_mask_prefix_sum.append(char_mask_prefix_sum[-1] + value)
-
-        loss_mask = []
-        for start, end in offset_mapping:
-            if end <= start:
-                loss_mask.append(0)
-            else:
-                loss_mask.append(1 if char_mask_prefix_sum[end] - char_mask_prefix_sum[start] > 0 else 0)
-
-        return token_ids, loss_mask
-
     def gen_multi_turn_loss_mask_distill_qwen(
         self, messages: list[dict], tools: list[dict] = None
     ) -> tuple[list[int], list[int]]:
@@ -221,8 +134,6 @@ def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> tuple
             return self.gen_multi_turn_loss_mask_qwen(messages, tools)
         elif self.tokenizer_type == "qwen3":
             return self.gen_multi_turn_loss_mask_qwen3(messages, tools)
-        elif self.tokenizer_type == "qwen3_5":
-            return self.gen_multi_turn_loss_mask_qwen3_5(messages, tools)
         elif self.tokenizer_type == "distill_qwen":
             return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools)
         else:
diff --git a/slime/utils/megatron_bridge_utils.py b/slime/utils/megatron_bridge_utils.py
index c87fb5b7b0..9e5f065cd4 100644
--- a/slime/utils/megatron_bridge_utils.py
+++ b/slime/utils/megatron_bridge_utils.py
@@ -6,38 +6,6 @@
     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/utils/memory_utils.py b/slime/utils/memory_utils.py
index d4b2d89321..c12f3cd0bc 100644
--- a/slime/utils/memory_utils.py
+++ b/slime/utils/memory_utils.py
@@ -1,7 +1,6 @@
 import gc
 import logging
 
-import psutil
 import torch
 import torch.distributed as dist
 
@@ -19,7 +18,6 @@ def clear_memory(clear_host_memory: bool = False):
 def available_memory():
     device = torch.cuda.current_device()
     free, total = torch.cuda.mem_get_info(device)
-    vm = psutil.virtual_memory()
     return {
         "gpu": str(device),
         "total_GB": _byte_to_gb(total),
@@ -27,10 +25,6 @@ def available_memory():
         "used_GB": _byte_to_gb(total - free),
         "allocated_GB": _byte_to_gb(torch.cuda.memory_allocated(device)),
         "reserved_GB": _byte_to_gb(torch.cuda.memory_reserved(device)),
-        "host_total_GB": _byte_to_gb(vm.total),
-        "host_available_GB": _byte_to_gb(vm.available),
-        "host_used_GB": _byte_to_gb(vm.used),
-        "host_free_GB": _byte_to_gb(vm.free),
     }
 
 
diff --git a/slime/utils/misc.py b/slime/utils/misc.py
index 4d4775f4f7..5a0841568a 100644
--- a/slime/utils/misc.py
+++ b/slime/utils/misc.py
@@ -119,6 +119,7 @@ def group_by(iterable, key=None):
     return dict(ret)
 
 
+# TODO fsdp can also use this
 def chunk_named_params_by_size(named_params: Iterable[tuple[str, torch.Tensor]], chunk_size: int):
     return _chunk_by_size(
         named_params,
diff --git a/slime/utils/ppo_utils.py b/slime/utils/ppo_utils.py
index 2a858e7a3f..2404883ab3 100644
--- a/slime/utils/ppo_utils.py
+++ b/slime/utils/ppo_utils.py
@@ -648,31 +648,30 @@ def chunked_gae(
 
 def calculate_log_probs_and_entropy(logits, tokens, tp_group, with_entropy: bool = False, chunk_size: int = -1):
     logits = logits.contiguous()
+    # TODO: not sure why we need to clone the logits here.
+    # Without the clone, the backward will trigger inplace edit error.
+    # It seems that the function with tp will modify the logits inplace.
     entropy = None
     if logits.size(0) != 0:
         if chunk_size > 0:
             num_chunks = (logits.size(0) - 1) // chunk_size + 1
-            logits_chunks = logits.chunk(num_chunks, dim=0)
             tokens_chunks = tokens.chunk(num_chunks, dim=0)
-
-            if with_entropy:
-                entropys = []
-                for logits_chunk in logits_chunks:
-                    entropy_input = logits_chunk.clone()
-                    entropys.append(compute_entropy_from_logits(entropy_input, tp_group))
-                entropy = torch.cat(entropys, dim=0)
-
+            logits_chunks = logits.chunk(num_chunks, dim=0)
             log_probs = []
             for tokens_chunk, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True):
                 log_prob = compute_log_probs(logits_chunk.clone(), tokens_chunk, tp_group)
                 log_probs.append(log_prob)
             log_prob = torch.cat(log_probs, dim=0)
-        else:
             if with_entropy:
-                entropy_input = logits.clone()
-                entropy = compute_entropy_from_logits(entropy_input, tp_group)
-
+                entropys = []
+                for _, logits_chunk in zip(tokens_chunks, logits_chunks, strict=True):
+                    entropy = compute_entropy_from_logits(logits_chunk.clone(), tp_group)
+                    entropys.append(entropy)
+                entropy = torch.cat(entropys, dim=0)
+        else:
             log_prob = compute_log_probs(logits.clone(), tokens, tp_group)
+            if with_entropy:
+                entropy = compute_entropy_from_logits(logits.clone(), tp_group)
     else:
         log_prob = logits.new_zeros((0,))
         if with_entropy:
diff --git a/slime/utils/processing_utils.py b/slime/utils/processing_utils.py
index fb652e73d0..18aa279971 100644
--- a/slime/utils/processing_utils.py
+++ b/slime/utils/processing_utils.py
@@ -1,10 +1,7 @@
 import base64
 import io
-import json
 import logging
-from pathlib import Path
 
-from PIL import Image
 from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizerBase, ProcessorMixin
 
 logger = logging.getLogger(__name__)
@@ -21,16 +18,17 @@ def load_tokenizer(name_or_path: str, **kwargs):
 
 def build_processor_kwargs(multimodal_inputs: dict | None = None) -> dict:
 
+    forced = {
+        # force return_tensors to None for input_ids
+        "return_tensors": None,
+    }
     modality_forced = {"return_tensors": "pt"}
 
     result = dict(multimodal_inputs) if multimodal_inputs else {}
 
-    # return_tensors=None for text (input_ids as lists), "pt" for modality-specific outputs
-    result["text_kwargs"] = {
-        **result.get("text_kwargs", {}),
-        "return_tensors": None,
-        "return_mm_token_type_ids": False,
-    }
+    result.update(forced)
+
+    # set return_tensors="pt" for modality-specific outputs
     for key in ("audio_kwargs", "images_kwargs", "videos_kwargs"):
         if key in result:
             result[key] = {**result[key], **modality_forced}
@@ -40,48 +38,6 @@ def build_processor_kwargs(multimodal_inputs: dict | None = None) -> dict:
     return result
 
 
-def _try_load_glm4v_processor(name_or_path: str, **kwargs):
-    """Fallback: manually construct a Glm4vProcessor for GLM-4.6V / GLM-4.5V models.
-
-    AutoProcessor fails for these models on transformers < 5.0 because
-    the Glm46VProcessor / Glm4vMoeProcessor classes are not registered.
-    The underlying Glm4vProcessor (non-MoE) works for both variants since
-    they share the same vision architecture.
-    """
-    try:
-        from transformers.models.glm4v.image_processing_glm4v import Glm4vImageProcessor
-        from transformers.models.glm4v.processing_glm4v import Glm4vProcessor
-        from transformers.models.glm4v.video_processing_glm4v import Glm4vVideoProcessor
-    except ImportError:
-        return None
-
-    pp_path = Path(name_or_path) / "preprocessor_config.json"
-    vp_path = Path(name_or_path) / "video_preprocessor_config.json"
-    if not pp_path.exists():
-        return None
-
-    skip_keys = {"image_processor_type", "processor_class", "video_processor_type"}
-    with open(pp_path) as f:
-        pp_cfg = {k: v for k, v in json.load(f).items() if k not in skip_keys}
-    image_processor = Glm4vImageProcessor(**pp_cfg)
-
-    video_processor = None
-    if vp_path.exists():
-        with open(vp_path) as f:
-            vp_cfg = {k: v for k, v in json.load(f).items() if k not in skip_keys}
-        video_processor = Glm4vVideoProcessor(**vp_cfg)
-
-    tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs)
-    proc = Glm4vProcessor(
-        image_processor=image_processor,
-        tokenizer=tokenizer,
-        video_processor=video_processor,
-        chat_template=tokenizer.chat_template,
-    )
-    logger.info(f"Loaded Glm4vProcessor manually for {name_or_path}")
-    return proc
-
-
 def load_processor(name_or_path: str, **kwargs):
     try:
         proc = AutoProcessor.from_pretrained(name_or_path, **kwargs)
@@ -89,67 +45,25 @@ def load_processor(name_or_path: str, **kwargs):
         logger.warning(f"Failed to load processor from {name_or_path}: {e}")
         proc = None
 
-    # If HF returned a tokenizer instead of a proper processor, discard it.
+    # If HF returned a tokenizer, discard it.
     if isinstance(proc, PreTrainedTokenizerBase) or not isinstance(proc, ProcessorMixin):
-        # Fallback: try to construct a GLM-4.6V / GLM-4.5V processor manually.
-        proc = _try_load_glm4v_processor(name_or_path, **kwargs)
+        proc = None
 
     return proc
 
 
-def _extract_images_from_messages(messages):
-    """Extract PIL images from chat messages containing multimodal content.
-
-    Handles base64 strings (with or without data: URI prefix), file paths,
-    and PIL Image objects embedded in message content dicts.
-    """
-    images = []
-    for msg in messages:
-        content = msg.get("content", [])
-        if not isinstance(content, list):
-            continue
-        for item in content:
-            if not isinstance(item, dict) or item.get("type") != "image":
-                continue
-            image_data = item.get("image")
-            if image_data is None:
-                continue
-            if isinstance(image_data, Image.Image):
-                images.append(image_data)
-            elif isinstance(image_data, str):
-                if image_data.startswith("data:"):
-                    _, encoded = image_data.split(",", 1)
-                    images.append(Image.open(io.BytesIO(base64.b64decode(encoded))))
-                else:
-                    try:
-                        raw = base64.b64decode(image_data)
-                        images.append(Image.open(io.BytesIO(raw)))
-                    except Exception:
-                        # Not base64 — try as file path
-                        images.append(Image.open(image_data))
-    return images
-
-
 def process_vision_info(prompt, processor):
-    """Extract PIL images (and videos) from the message list for training.
-
-    Tries qwen_vl_utils first (Qwen VL family), falls back to generic
-    extraction for other models (e.g. GLM-4.6V).
-    """
-    try:
-        from qwen_vl_utils import process_vision_info as qwen_process_vision_info
-
-        if hasattr(processor.image_processor, "patch_size"):
-            image_patch_size = processor.image_processor.patch_size
-        else:
-            image_patch_size = DEFAULT_PATCH_SIZE
-        images, videos = qwen_process_vision_info(prompt, image_patch_size=image_patch_size)
-    except Exception:
-        # Fallback: generic extraction for non-Qwen models
-        images = _extract_images_from_messages(prompt) or None
-        videos = None
-
-    return {"images": images, "videos": videos}
+    # TODO: temporary solution, will write image utils for slime later
+    from qwen_vl_utils import process_vision_info as qwen_process_vision_info
+
+    if hasattr(processor.image_processor, "patch_size"):
+        image_patch_size = processor.image_processor.patch_size
+    else:
+        logger.info(f"Using default patch size: {DEFAULT_PATCH_SIZE}")
+        image_patch_size = DEFAULT_PATCH_SIZE
+    images, videos = qwen_process_vision_info(prompt, image_patch_size=image_patch_size)
+    multimodal_inputs = {"images": images, "videos": videos}
+    return multimodal_inputs
 
 
 def encode_image_for_rollout_engine(image) -> str:
diff --git a/slime/utils/reloadable_process_group.py b/slime/utils/reloadable_process_group.py
index 66055a352f..a8ea33a909 100644
--- a/slime/utils/reloadable_process_group.py
+++ b/slime/utils/reloadable_process_group.py
@@ -296,7 +296,7 @@ def reload_process_groups():
 def _wrap_low_level_call():
     try:
         mem_info = available_memory()
-        if mem_info["free_GB"] < 3:
+        if mem_info["free_GB"] < 5:
             clear_memory()
         yield
     except Exception as e:
diff --git a/slime/utils/trace_utils.py b/slime/utils/trace_utils.py
deleted file mode 100644
index b51bc5b931..0000000000
--- a/slime/utils/trace_utils.py
+++ /dev/null
@@ -1,612 +0,0 @@
-from __future__ import annotations
-
-import contextvars
-import functools
-import inspect
-import logging
-import time
-import uuid
-from collections.abc import Callable
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-from typing import Any
-
-from slime.utils.types import Sample
-
-TRACE_VERSION = 1
-SGLANG_TRACE_META_KEYS = (
-    "prompt_tokens",
-    "completion_tokens",
-    "cached_tokens",
-    "pd_prefill_bootstrap_queue_duration",
-    "pd_prefill_forward_duration",
-    "pd_prefill_transfer_queue_duration",
-    "pd_prefill_retry_count",
-    "pd_decode_prealloc_duration",
-    "pd_decode_transfer_duration",
-    "pd_decode_forward_duration",
-    "pd_bootstrap_duration",
-    "pd_alloc_waiting_duration",
-    "pd_transfer_speed_gb_s",
-    "pd_transfer_total_mb",
-)
-
-logger = logging.getLogger(__name__)
-_TRACE_STACK: contextvars.ContextVar[tuple[tuple[str, str], ...]] = contextvars.ContextVar(
-    "slime_trace_stack",
-    default=(),
-)
-_TRACE_HANDLE_STACK: contextvars.ContextVar[tuple[tuple[TraceHandle, ...], ...]] = contextvars.ContextVar(
-    "slime_trace_handle_stack",
-    default=(),
-)
-_TRACE_AUTO_INFER_WARNED: set[str] = set()
-
-
-@dataclass
-class TraceHandle:
-    trace_id: str
-    carrier: dict[str, Any]
-    sample_id: int | str | None = None
-    group_id: int | str | None = None
-    attempt: int = 0
-    parent_span_id: str | None = None
-
-
-@dataclass
-class TraceSpanContext:
-    target: Sample | TraceHandle | list[Sample | TraceHandle]
-    handles: list[TraceHandle]
-    end_attrs: dict[str, Any] = field(default_factory=dict)
-    end_events: list[dict[str, Any]] = field(default_factory=list)
-    closed: bool = False
-
-    def set(self, key: str, value: Any) -> TraceSpanContext:
-        self.end_attrs[key] = value
-        self._sync_end_events({key: value})
-        return self
-
-    def update(self, attrs: dict[str, Any] | None) -> TraceSpanContext:
-        if attrs:
-            self.end_attrs.update(attrs)
-            self._sync_end_events(attrs)
-        return self
-
-    def set_attr(self, key: str, value: Any) -> TraceSpanContext:
-        return self.set(key, value)
-
-    def update_attrs(self, attrs: dict[str, Any] | None) -> TraceSpanContext:
-        return self.update(attrs)
-
-    def build_end_attrs(self) -> dict[str, Any] | None:
-        return dict(self.end_attrs) or None
-
-    def finalize(self, end_events: list[dict[str, Any]]) -> None:
-        self.end_events = end_events
-        self.closed = True
-        if self.end_attrs:
-            self._sync_end_events(self.end_attrs)
-
-    def _sync_end_events(self, attrs: dict[str, Any]) -> None:
-        if not self.end_events or not attrs:
-            return
-        for event in self.end_events:
-            event.setdefault("attrs", {})
-            event["attrs"].update(attrs)
-
-
-def _noop_handle() -> TraceHandle:
-    return TraceHandle(
-        trace_id="",
-        carrier={
-            "version": TRACE_VERSION,
-            "trace_id": "",
-            "events": [],
-            "sample_id": None,
-            "group_id": None,
-            "attempt": 0,
-        },
-    )
-
-
-def _log_trace_error(action: str, exc: Exception) -> None:
-    logger.debug("trace %s skipped: %s", action, exc, exc_info=True)
-
-
-def _new_trace_id() -> str:
-    return uuid.uuid4().hex
-
-
-def _new_span_id() -> str:
-    return uuid.uuid4().hex
-
-
-def build_sglang_meta_trace_attrs(meta: dict[str, Any]) -> dict[str, Any]:
-    attrs = {key: meta[key] for key in SGLANG_TRACE_META_KEYS if key in meta and meta[key] is not None}
-    attrs["finish_reason"] = meta["finish_reason"]["type"]
-    return attrs
-
-
-def _ensure_trace_carrier(
-    carrier: dict[str, Any] | None,
-    *,
-    trace_id: str | None = None,
-    sample_id: int | str | None = None,
-    group_id: int | str | None = None,
-    attempt: int = 0,
-) -> dict[str, Any]:
-    if carrier is None:
-        carrier = {}
-    carrier.setdefault("version", TRACE_VERSION)
-    carrier.setdefault("trace_id", trace_id or _new_trace_id())
-    carrier.setdefault("events", [])
-    if sample_id is not None:
-        carrier["sample_id"] = sample_id
-    else:
-        carrier.setdefault("sample_id", None)
-    if group_id is not None:
-        carrier["group_id"] = group_id
-    else:
-        carrier.setdefault("group_id", None)
-    carrier["attempt"] = int(carrier.get("attempt", attempt))
-    return carrier
-
-
-def bind_trace(sample: Sample) -> TraceHandle:
-    try:
-        sample.trace = _ensure_trace_carrier(
-            getattr(sample, "trace", None),
-            sample_id=sample.index,
-            group_id=sample.group_index,
-        )
-        return TraceHandle(
-            trace_id=sample.trace["trace_id"],
-            carrier=sample.trace,
-            sample_id=sample.trace.get("sample_id"),
-            group_id=sample.trace.get("group_id"),
-            attempt=int(sample.trace.get("attempt", 0)),
-        )
-    except Exception as exc:
-        _log_trace_error("bind", exc)
-        return _noop_handle()
-
-
-def bind_trace_carrier(
-    carrier: dict[str, Any] | None,
-    *,
-    trace_id: str | None = None,
-    sample_id: int | str | None = None,
-    group_id: int | str | None = None,
-    attempt: int = 0,
-    parent_span_id: str | None = None,
-) -> TraceHandle:
-    try:
-        trace = _ensure_trace_carrier(
-            carrier,
-            trace_id=trace_id,
-            sample_id=sample_id,
-            group_id=group_id,
-            attempt=attempt,
-        )
-        return TraceHandle(
-            trace_id=trace["trace_id"],
-            carrier=trace,
-            sample_id=trace.get("sample_id"),
-            group_id=trace.get("group_id"),
-            attempt=int(trace.get("attempt", 0)),
-            parent_span_id=parent_span_id,
-        )
-    except Exception as exc:
-        _log_trace_error("bind_carrier", exc)
-        handle = _noop_handle()
-        handle.parent_span_id = parent_span_id
-        return handle
-
-
-def export_trace(handle: TraceHandle) -> dict[str, Any]:
-    try:
-        return {
-            "version": TRACE_VERSION,
-            "trace_id": handle.trace_id,
-            "sample_id": handle.sample_id,
-            "group_id": handle.group_id,
-            "attempt": handle.attempt,
-            "parent_span_id": handle.parent_span_id or _get_current_parent_span_id(handle.trace_id),
-        }
-    except Exception as exc:
-        _log_trace_error("export", exc)
-        return {
-            "version": TRACE_VERSION,
-            "trace_id": "",
-            "sample_id": None,
-            "group_id": None,
-            "attempt": 0,
-            "parent_span_id": None,
-        }
-
-
-def import_trace(payload: dict[str, Any], carrier: dict[str, Any] | None = None) -> TraceHandle:
-    try:
-        return bind_trace_carrier(
-            carrier,
-            trace_id=payload.get("trace_id"),
-            sample_id=payload.get("sample_id"),
-            group_id=payload.get("group_id"),
-            attempt=int(payload.get("attempt", 0)),
-            parent_span_id=payload.get("parent_span_id"),
-        )
-    except Exception as exc:
-        _log_trace_error("import", exc)
-        return _noop_handle()
-
-
-def trace_event(
-    target: Sample | TraceHandle | list[Sample | TraceHandle], name: str, *, attrs: dict[str, Any] | None = None
-):
-    try:
-        timestamp = time.time()
-        for handle in _coerce_handles(target):
-            _append_event(handle, kind="event", name=name, timestamp=timestamp, attrs=attrs)
-    except Exception as exc:
-        _log_trace_error(f"event:{name}", exc)
-
-
-@contextmanager
-def trace_span(
-    target: Sample | TraceHandle | list[Sample | TraceHandle],
-    name: str,
-    *,
-    attrs: dict[str, Any] | None = None,
-    record_error: bool = True,
-):
-    try:
-        handles = _coerce_handles(target)
-    except Exception as exc:
-        _log_trace_error(f"span:{name}", exc)
-        handles = []
-
-    if not handles:
-        yield target
-        return
-
-    timestamp = time.time()
-    stack_before = _TRACE_STACK.get()
-    handle_stack_before = _TRACE_HANDLE_STACK.get()
-    span_records: list[tuple[TraceHandle, str]] = []
-    new_entries: list[tuple[str, str]] = []
-
-    try:
-        for handle in handles:
-            span_id = _new_span_id()
-            parent_span_id = handle.parent_span_id or _get_current_parent_span_id(handle.trace_id, stack=stack_before)
-            _append_event(
-                handle,
-                kind="span_start",
-                name=name,
-                timestamp=timestamp,
-                span_id=span_id,
-                parent_span_id=parent_span_id,
-                attrs=attrs,
-            )
-            span_records.append((handle, span_id))
-            new_entries.append((handle.trace_id, span_id))
-        token = _TRACE_STACK.set(stack_before + tuple(new_entries))
-        handle_token = _TRACE_HANDLE_STACK.set(handle_stack_before + (tuple(handles),))
-    except Exception as exc:
-        _log_trace_error(f"span:{name}", exc)
-        yield target
-        return
-
-    span_context = TraceSpanContext(
-        target=handles[0] if len(handles) == 1 else handles,
-        handles=handles,
-    )
-
-    try:
-        yield span_context
-    except Exception as exc:
-        try:
-            end_attrs = span_context.build_end_attrs()
-            if record_error:
-                error_attrs = {"error_type": type(exc).__name__, "error_message": str(exc)}
-                if end_attrs:
-                    end_attrs.update(error_attrs)
-                else:
-                    end_attrs = error_attrs
-            span_context.finalize(_record_span_end(span_records, name=name, attrs=end_attrs))
-        except Exception as trace_exc:
-            _log_trace_error(f"span_end:{name}", trace_exc)
-        raise
-    else:
-        try:
-            span_context.finalize(_record_span_end(span_records, name=name, attrs=span_context.build_end_attrs()))
-        except Exception as exc:
-            _log_trace_error(f"span_end:{name}", exc)
-    finally:
-        try:
-            _TRACE_STACK.reset(token)
-        except Exception as exc:
-            _log_trace_error(f"span_reset:{name}", exc)
-        try:
-            _TRACE_HANDLE_STACK.reset(handle_token)
-        except Exception as exc:
-            _log_trace_error(f"span_handle_reset:{name}", exc)
-
-
-def trace_next_attempt(
-    target: Sample | TraceHandle | list[Sample | TraceHandle],
-    *,
-    attrs: dict[str, Any] | None = None,
-):
-    try:
-        handles = _coerce_handles(target)
-        for handle in handles:
-            next_attempt = int(handle.carrier.get("attempt", 0)) + 1
-            handle.carrier["attempt"] = next_attempt
-            handle.attempt = next_attempt
-            attempt_attrs = {"attempt": next_attempt}
-            if attrs:
-                attempt_attrs.update(attrs)
-            trace_event(handle, "attempt_start", attrs=attempt_attrs)
-        if not handles:
-            return target
-        return handles[0] if len(handles) == 1 else handles
-    except Exception as exc:
-        _log_trace_error("next_attempt", exc)
-        return target
-
-
-def trace_function(
-    name: str,
-    *,
-    target: str | None = None,
-    target_getter: Callable[..., Sample | TraceHandle | list[Sample | TraceHandle] | None] | None = None,
-    attrs_getter: Callable[..., dict[str, Any] | None] | None = None,
-    record_error: bool = True,
-):
-    def decorator(fn):
-        if inspect.iscoroutinefunction(fn):
-
-            @functools.wraps(fn)
-            async def async_wrapper(*args, **kwargs):
-                resolved_target = _resolve_trace_function_target(
-                    fn,
-                    args,
-                    kwargs,
-                    target=target,
-                    target_getter=target_getter,
-                )
-                if resolved_target is None:
-                    return await fn(*args, **kwargs)
-                attrs = _resolve_trace_function_attrs(fn, args, kwargs, attrs_getter=attrs_getter)
-                with trace_span(resolved_target, name, attrs=attrs, record_error=record_error):
-                    return await fn(*args, **kwargs)
-
-            return async_wrapper
-
-        @functools.wraps(fn)
-        def sync_wrapper(*args, **kwargs):
-            resolved_target = _resolve_trace_function_target(
-                fn,
-                args,
-                kwargs,
-                target=target,
-                target_getter=target_getter,
-            )
-            if resolved_target is None:
-                return fn(*args, **kwargs)
-            attrs = _resolve_trace_function_attrs(fn, args, kwargs, attrs_getter=attrs_getter)
-            with trace_span(resolved_target, name, attrs=attrs, record_error=record_error):
-                return fn(*args, **kwargs)
-
-        return sync_wrapper
-
-    return decorator
-
-
-def _record_span_end(
-    span_records: list[tuple[TraceHandle, str]],
-    *,
-    name: str,
-    attrs: dict[str, Any] | None,
-) -> list[dict[str, Any]]:
-    timestamp = time.time()
-    events = []
-    for handle, span_id in span_records:
-        events.append(
-            _append_event(
-                handle,
-                kind="span_end",
-                name=name,
-                timestamp=timestamp,
-                span_id=span_id,
-                attrs=attrs,
-            )
-        )
-    return events
-
-
-def _append_event(
-    handle: TraceHandle,
-    *,
-    kind: str,
-    name: str,
-    timestamp: float,
-    attrs: dict[str, Any] | None = None,
-    span_id: str | None = None,
-    parent_span_id: str | None = None,
-) -> dict[str, Any]:
-    event = {
-        "type": kind,
-        "name": name,
-        "ts": timestamp,
-        "trace_id": handle.trace_id,
-        "sample_id": handle.sample_id,
-        "group_id": handle.group_id,
-        "attempt": int(handle.carrier.get("attempt", handle.attempt)),
-    }
-    if span_id is not None:
-        event["span_id"] = span_id
-    if parent_span_id is not None:
-        event["parent_span_id"] = parent_span_id
-    if attrs:
-        event["attrs"] = dict(attrs)
-    handle.carrier["events"].append(event)
-    return event
-
-
-def _coerce_handles(target: Sample | TraceHandle | list[Sample | TraceHandle]) -> list[TraceHandle]:
-    target = _adapt_trace_target(target)
-    if isinstance(target, TraceHandle):
-        return [target]
-    if isinstance(target, Sample):
-        return [bind_trace(target)]
-    if isinstance(target, list):
-        handles = []
-        for item in target:
-            handles.extend(_coerce_handles(item))
-        return handles
-    return []
-
-
-def _get_current_parent_span_id(
-    trace_id: str,
-    *,
-    stack: tuple[tuple[str, str], ...] | None = None,
-) -> str | None:
-    stack = _TRACE_STACK.get() if stack is None else stack
-    for current_trace_id, span_id in reversed(stack):
-        if current_trace_id == trace_id:
-            return span_id
-    return None
-
-
-def _resolve_trace_function_target(
-    fn,
-    args: tuple[Any, ...],
-    kwargs: dict[str, Any],
-    *,
-    target: str | None,
-    target_getter: Callable[..., Sample | TraceHandle | list[Sample | TraceHandle] | None] | None,
-):
-    try:
-        bound = inspect.signature(fn).bind_partial(*args, **kwargs)
-    except Exception as exc:
-        _log_trace_error(f"trace_function_bind:{getattr(fn, '__qualname__', fn)}", exc)
-        bound = None
-
-    if target is not None:
-        if bound is None or target not in bound.arguments:
-            logger.warning(
-                "trace_function target '%s' not found for %s; tracing disabled for this call",
-                target,
-                getattr(fn, "__qualname__", repr(fn)),
-            )
-            return None
-        resolved = _normalize_trace_target(bound.arguments.get(target))
-        if resolved is None:
-            logger.warning(
-                "trace_function target '%s' for %s is not a supported trace target; tracing disabled for this call",
-                target,
-                getattr(fn, "__qualname__", repr(fn)),
-            )
-        return resolved
-
-    if target_getter is not None:
-        try:
-            resolved = _normalize_trace_target(target_getter(*args, **kwargs))
-            return resolved
-        except Exception as exc:
-            _log_trace_error(f"trace_function_target_getter:{getattr(fn, '__qualname__', fn)}", exc)
-            return None
-
-    inferred = _infer_trace_target(bound.arguments.values() if bound is not None else args)
-    if inferred is not None:
-        warn_key = getattr(fn, "__module__", "") + "." + getattr(fn, "__qualname__", repr(fn))
-        if warn_key not in _TRACE_AUTO_INFER_WARNED:
-            _TRACE_AUTO_INFER_WARNED.add(warn_key)
-            logger.warning(
-                "trace_function auto-inferred target for %s; inference may be ambiguous, prefer explicit target=...",
-                getattr(fn, "__qualname__", repr(fn)),
-            )
-        return inferred
-
-    return _get_current_trace_target()
-
-
-def _resolve_trace_function_attrs(
-    fn,
-    args: tuple[Any, ...],
-    kwargs: dict[str, Any],
-    *,
-    attrs_getter: Callable[..., dict[str, Any] | None] | None,
-) -> dict[str, Any] | None:
-    if attrs_getter is None:
-        return None
-    try:
-        attrs = attrs_getter(*args, **kwargs)
-        if attrs is None:
-            return None
-        if isinstance(attrs, dict):
-            return attrs
-        logger.warning(
-            "trace_function attrs_getter for %s returned non-dict %s; ignoring attrs",
-            getattr(fn, "__qualname__", repr(fn)),
-            type(attrs).__name__,
-        )
-        return None
-    except Exception as exc:
-        _log_trace_error(f"trace_function_attrs_getter:{getattr(fn, '__qualname__', fn)}", exc)
-        return None
-
-
-def _infer_trace_target(values) -> Sample | TraceHandle | list[Sample | TraceHandle] | None:
-    for value in values:
-        normalized = _normalize_trace_target(value)
-        if normalized is not None:
-            return normalized
-    return None
-
-
-def _normalize_trace_target(value):
-    value = _adapt_trace_target(value)
-    if isinstance(value, (Sample, TraceHandle)):
-        return value
-    if isinstance(value, list) and value:
-        if all(_normalize_trace_target(item) is not None for item in value):
-            return value
-    return None
-
-
-def _adapt_trace_target(value):
-    if value is None:
-        return None
-    if isinstance(value, (Sample, TraceHandle)):
-        return value
-    if isinstance(value, list):
-        return [_adapt_trace_target(item) for item in value]
-    if _looks_like_sample_box(value):
-        generation = getattr(value, "generation", None)
-        if generation:
-            return generation
-        return getattr(value, "prompt_sample", None)
-    return value
-
-
-def _get_current_trace_target() -> TraceHandle | list[TraceHandle] | None:
-    handle_stack = _TRACE_HANDLE_STACK.get()
-    if not handle_stack:
-        return None
-    current_handles = list(handle_stack[-1])
-    if not current_handles:
-        return None
-    if len(current_handles) == 1:
-        return current_handles[0]
-    return current_handles
-
-
-def _looks_like_sample_box(value: Any) -> bool:
-    cls = getattr(value, "__class__", None)
-    if cls is None or getattr(cls, "__name__", "") != "SampleBox":
-        return False
-    return hasattr(value, "prompt_sample") and hasattr(value, "generation")
diff --git a/slime/utils/wandb_utils.py b/slime/utils/wandb_utils.py
index 2ec859de5f..bb491d1b03 100644
--- a/slime/utils/wandb_utils.py
+++ b/slime/utils/wandb_utils.py
@@ -79,62 +79,6 @@ def init_wandb_primary(args):
     args.wandb_run_id = wandb.run.id
 
 
-def reinit_wandb_primary_with_open_metrics(args, router_addr):
-    """Re-initialize the primary W&B run with open metrics endpoints.
-
-    The primary wandb init happens before rollout servers start (to obtain
-    ``wandb_run_id`` for secondary processes).  This function is called
-    *after* servers are up so the router address is available for scraping
-    SGLang Prometheus metrics via the primary process's stats monitor.
-    """
-    if not args.use_wandb or _is_offline_mode(args):
-        return
-    if getattr(args, "wandb_mode", None) == "disabled":
-        return
-    if router_addr is None:
-        return
-    wandb_run_id = getattr(args, "wandb_run_id", None)
-    if wandb_run_id is None:
-        return
-
-    import sglang_router
-
-    if "slime" not in sglang_router.__version__:
-        logger.warning(
-            "Only customized sglang_router from https://github.com/zhuzilin/sgl-router supports uploading metrics."
-        )
-        return
-
-    logger.info(f"Re-initializing primary W&B with SGLang metrics at {router_addr}.")
-
-    wandb.finish()
-
-    init_kwargs = {
-        "id": wandb_run_id,
-        "entity": args.wandb_team,
-        "project": args.wandb_project,
-        "resume": "allow",
-        "reinit": True,
-        "settings": wandb.Settings(
-            mode="shared",
-            x_primary=True,
-            x_stats_open_metrics_endpoints={
-                "sgl_engine": f"{router_addr}/engine_metrics",
-            },
-            x_stats_open_metrics_filters={
-                "sgl_engine.*": {},
-            },
-        ),
-    }
-
-    if args.wandb_dir:
-        os.makedirs(args.wandb_dir, exist_ok=True)
-        init_kwargs["dir"] = args.wandb_dir
-
-    wandb.init(**init_kwargs)
-    _init_wandb_common()
-
-
 def _compute_config_for_logging(args):
     output = deepcopy(args.__dict__)
 
@@ -148,7 +92,7 @@ def _compute_config_for_logging(args):
 
 
 # https://docs.wandb.ai/guides/track/log/distributed-training/#track-all-processes-to-a-single-run
-def init_wandb_secondary(args):
+def init_wandb_secondary(args, router_addr=None):
     wandb_run_id = getattr(args, "wandb_run_id", None)
     if wandb_run_id is None:
         return
@@ -172,6 +116,17 @@ def init_wandb_secondary(args):
             x_update_finish_state=False,
         )
 
+    if getattr(args, "sglang_enable_metrics", False) and router_addr is not None:
+        logger.info(f"Forward SGLang metrics at {router_addr} to WandB.")
+        settings_kwargs |= dict(
+            x_stats_open_metrics_endpoints={
+                "sgl_engine": f"{router_addr}/engine_metrics",
+            },
+            x_stats_open_metrics_filters={
+                "sgl_engine.*": {},
+            },
+        )
+
     init_kwargs = {
         "id": wandb_run_id,
         "entity": args.wandb_team,
diff --git a/slime_plugins/mbridge/deepseek_v32.py b/slime_plugins/mbridge/deepseek_v32.py
index d45fd40fe6..d3a4ee7977 100644
--- a/slime_plugins/mbridge/deepseek_v32.py
+++ b/slime_plugins/mbridge/deepseek_v32.py
@@ -6,15 +6,6 @@
 
 @register_model("deepseek_v32")
 class DeepseekV32Bridge(DeepseekV3Bridge):
-
-    def __init__(self, hf_config, **kwargs):
-        # transformers 5.x stores rope_theta inside rope_parameters dict,
-        # but DeepseekV3Bridge._build_config() expects hf_config.rope_theta directly.
-        if not hasattr(hf_config, "rope_theta"):
-            rope_params = getattr(hf_config, "rope_parameters", None) or {}
-            hf_config.rope_theta = rope_params.get("rope_theta", 1000000)
-        super().__init__(hf_config, **kwargs)
-
     _DSA_ATTENTION_MAPPING = {
         "self_attention.wq_b.weight": ["model.layers.{layer_number}.self_attn.indexer.wq_b.weight"],
         "self_attention.wk.weight": ["model.layers.{layer_number}.self_attn.indexer.wk.weight"],
diff --git a/slime_plugins/mbridge/glm4moe_lite.py b/slime_plugins/mbridge/glm4moe_lite.py
index ceaa8d86e0..8306d281ae 100644
--- a/slime_plugins/mbridge/glm4moe_lite.py
+++ b/slime_plugins/mbridge/glm4moe_lite.py
@@ -1,76 +1,7 @@
-import torch
 from mbridge.core import register_model
-from mbridge.core.safetensor_io import SafeTensorIO
 from mbridge.models import DeepseekV3Bridge
 
 
 @register_model("glm4_moe_lite")
 class GLM4MoELiteBridge(DeepseekV3Bridge):
-    """
-    Bridge for GLM-4.7-Flash (glm4_moe_lite) models.
-
-    Extends DeepseekV3Bridge with:
-    - Dynamic MTP layer indexing (parent hardcodes layer 61 for DeepSeek V3)
-    - Standard bf16 safetensor loading (parent uses FP8 dequant for DeepSeek V3)
-    """
-
-    def __init__(self, hf_config, **kwargs):
-        # Patch rope_theta: GLM-4.7-Flash stores it in rope_parameters dict,
-        # but DeepseekV3Bridge._build_config() expects hf_config.rope_theta directly.
-        if not hasattr(hf_config, "rope_theta"):
-            rope_params = getattr(hf_config, "rope_parameters", None) or {}
-            hf_config.rope_theta = rope_params.get("rope_theta", 1000000)
-        super().__init__(hf_config, **kwargs)
-        # Override the shared state dict mapping with dynamic layer index.
-        # DeepseekV3Bridge hardcodes layer 61; GLM-4.7-Flash uses num_hidden_layers (47).
-        n = hf_config.num_hidden_layers
-        if getattr(hf_config, "num_nextn_predict_layers", 0) and n:
-            self._SHARED_STATE_DICT_MAPPING = {
-                "embedding.word_embeddings.weight": [
-                    "model.embed_tokens.weight",
-                    f"model.layers.{n}.embed_tokens.weight",
-                ],
-                "output_layer.weight": [
-                    "lm_head.weight",
-                    f"model.layers.{n}.shared_head.head.weight",
-                ],
-            }
-
-    def _get_safetensor_io(self, weights_path: str):
-        """Use standard SafeTensorIO — GLM-4.7-Flash ships bf16 safetensors, not FP8."""
-        return SafeTensorIO(self._get_actual_hf_path(weights_path))
-
-    def _weight_to_hf_format(
-        self, mcore_weights_name: str, mcore_weights: torch.Tensor
-    ) -> tuple[list[str], list[torch.Tensor]]:
-        """Handle shared embedding/output weights for MTP with dynamic layer count."""
-        if (
-            self.config.mtp_num_layers is not None
-            and self.config.mtp_num_layers >= 1
-            and mcore_weights_name in self._SHARED_STATE_DICT_MAPPING
-        ):
-            hf_names = self._SHARED_STATE_DICT_MAPPING[mcore_weights_name]
-            return hf_names, [mcore_weights] * len(hf_names)
-        # Skip DeepseekV3Bridge's _weight_to_hf_format (hardcoded 61) and go to Bridge base
-        return super(DeepseekV3Bridge, self)._weight_to_hf_format(mcore_weights_name, mcore_weights)
-
-    def _convert_mtp_param(self, name: str) -> list[str]:
-        """Convert MTP parameter names with dynamic layer count (not hardcoded 61)."""
-        assert self.config.mtp_num_layers == 1, "only support one mtp layer for now"
-        n = self.config.num_layers
-        direct_name_mapping = {
-            "mtp.layers.0.enorm.weight": f"model.layers.{n}.enorm.weight",
-            "mtp.layers.0.hnorm.weight": f"model.layers.{n}.hnorm.weight",
-            "mtp.layers.0.eh_proj.weight": f"model.layers.{n}.eh_proj.weight",
-            "mtp.layers.0.final_layernorm.weight": f"model.layers.{n}.shared_head.norm.weight",
-        }
-        if name in direct_name_mapping:
-            return [direct_name_mapping[name]]
-        assert "mtp.layers.0.transformer_layer" in name, f"mtp not found in {name}"
-        proxy_name = name.replace("mtp.layers.0.transformer_layer", f"decoder.layers.{n}")
-        if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name:
-            return self._weight_name_mapping_attention(proxy_name)
-        elif "mlp" in proxy_name:
-            return self._weight_name_mapping_mlp(proxy_name)
-        else:
-            raise NotImplementedError(f"Unsupported MTP parameter name: {name}")
+    pass
diff --git a/slime_plugins/mbridge/qwen3_5.py b/slime_plugins/mbridge/qwen3_5.py
index d01094ecc9..670227888b 100644
--- a/slime_plugins/mbridge/qwen3_5.py
+++ b/slime_plugins/mbridge/qwen3_5.py
@@ -1,5 +1,3 @@
-import inspect
-
 import torch
 from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec
 
@@ -117,44 +115,12 @@ 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)
-        if transformer_config_class is None:
-            return True
-
-        dataclass_fields = getattr(transformer_config_class, "__dataclass_fields__", None)
-        if dataclass_fields is not None:
-            return kwarg_name in dataclass_fields
-
-        try:
-            signature = inspect.signature(transformer_config_class)
-        except (TypeError, ValueError):
-            return True
-        return kwarg_name in signature.parameters
-
-    def _get_transformer_layer_spec(self, vp_stage=None):
-        transformer_layer_spec = super()._get_transformer_layer_spec(vp_stage)
-        self._last_transformer_layer_spec = transformer_layer_spec
-        return transformer_layer_spec
-
     def _get_gptmodel_args(self) -> dict:
         """Override to add MTP block spec if needed."""
         ret = super()._get_gptmodel_args()
         text_config = self._get_text_config()
         if getattr(text_config, "mtp_num_hidden_layers", None) is not None:
-            transformer_layer_spec = getattr(self, "_last_transformer_layer_spec", None)
-            if transformer_layer_spec is None:
-                transformer_layer_spec = self._get_transformer_layer_spec()
+            transformer_layer_spec = self.config
             mtp_block_spec = get_gpt_mtp_block_spec(self.config, transformer_layer_spec, use_transformer_engine=True)
             ret["mtp_block_spec"] = mtp_block_spec
         return ret
@@ -180,25 +146,6 @@ def _weight_name_mapping_mlp(self, name: str) -> list[str]:
             raise NotImplementedError(f"Unsupported parameter name: {name}")
         return convert_names
 
-    def _weight_name_mapping_mtp_mlp(self, name: str) -> list[str]:
-        """Handle MTP MLP mappings, keeping per-expert tensors unfused for MoE layers."""
-        layer_number = name.split(".")[2]
-        mapping = self._MTP_MLP_MAPPING if "mlp.experts.linear_fc" in name else self._MLP_MAPPING
-        convert_names = []
-        for keyword, mapping_names in mapping.items():
-            if keyword in name:
-                if "{expert_id}" in mapping_names[0]:
-                    expert_id = name.split("weight")[-1]
-                    convert_names.extend(
-                        [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names]
-                    )
-                else:
-                    convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names])
-                break
-        if len(convert_names) == 0:
-            raise NotImplementedError(f"Unsupported parameter name: {name}")
-        return convert_names
-
     def _weight_name_mapping_mcore_to_hf(self, mcore_weights_name: str) -> list[str]:
         """Override to handle MTP layer mappings."""
         if "mtp" in mcore_weights_name:
@@ -232,7 +179,7 @@ def _convert_mtp_param(self, name: str) -> list[str]:
             if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name:
                 convert_names = super()._weight_name_mapping_attention(proxy_name)
             elif "mlp" in proxy_name or "pre_mlp_layernorm" in proxy_name:
-                convert_names = self._weight_name_mapping_mtp_mlp(proxy_name)
+                convert_names = super()._weight_name_mapping_mlp(proxy_name)
             else:
                 raise NotImplementedError(f"Unsupported transformer component in MTP: {name}")
 
@@ -286,28 +233,23 @@ def _weight_to_mcore_format(
         if "mlp.experts.linear_fc" in mcore_weights_name and len(hf_weights) == 1:
             w = hf_weights[0]
             if w.dim() == 3:
-                # Extract local expert_id from name like "...linear_fc1.weight42"
-                local_expert_id = int(mcore_weights_name.split("weight")[-1])
-                # When using Expert Parallelism (EP), the local expert_id is relative
-                # to this EP rank. We need to convert to global expert_id to index
-                # into the full HF fused tensor [num_experts, ...].
-                from megatron.core import mpu
-
-                ep_size = mpu.get_expert_model_parallel_world_size()
-                if ep_size > 1:
-                    ep_rank = mpu.get_expert_model_parallel_rank()
-                    num_local_experts = w.shape[0] // ep_size
-                    global_expert_id = ep_rank * num_local_experts + local_expert_id
-                else:
-                    global_expert_id = local_expert_id
-                expert_w = w[global_expert_id]  # (out_features, in_features)
+                # Extract expert_id from name like "...linear_fc1.weight42"
+                expert_id = int(mcore_weights_name.split("weight")[-1])
+                expert_w = w[expert_id]  # (out_features, in_features)
                 return expert_w.contiguous()
 
-        return super()._weight_to_mcore_format(mcore_weights_name, hf_weights)
+        weight = super()._weight_to_mcore_format(mcore_weights_name, hf_weights)
+        if mcore_weights_name.endswith("eh_proj.weight"):
+            first_half, second_half = weight.chunk(2, dim=1)
+            weight = torch.cat([second_half, first_half], dim=1)
+        return weight
 
     def _weight_to_hf_format(
         self, mcore_weights_name: str, mcore_weights: torch.Tensor
     ) -> tuple[list[str], list[torch.Tensor]]:
+        if mcore_weights_name.endswith("eh_proj.weight"):
+            first_half, second_half = mcore_weights.chunk(2, dim=1)
+            mcore_weights = torch.cat([second_half, first_half], dim=1)
         return super()._weight_to_hf_format(mcore_weights_name, mcore_weights)
 
     def _build_config(self):
@@ -331,9 +273,6 @@ def _build_config(self):
             **mtp_args,
         )
 
-        if self._supports_transformer_config_kwarg("use_gated_attention"):
-            base_kwargs["use_gated_attention"] = True
-
         # Handle MoE-specific config
         if hasattr(text_config, "num_experts"):
             base_kwargs.update(
diff --git a/slime_plugins/megatron_bridge/__init__.py b/slime_plugins/megatron_bridge/__init__.py
index a0425d491b..e69de29bb2 100644
--- a/slime_plugins/megatron_bridge/__init__.py
+++ b/slime_plugins/megatron_bridge/__init__.py
@@ -1 +0,0 @@
-import slime_plugins.megatron_bridge.glm4v_moe  # noqa: F401  # register GLM-4.6V bridge
diff --git a/slime_plugins/megatron_bridge/glm4v_moe.py b/slime_plugins/megatron_bridge/glm4v_moe.py
deleted file mode 100644
index 785276e976..0000000000
--- a/slime_plugins/megatron_bridge/glm4v_moe.py
+++ /dev/null
@@ -1,711 +0,0 @@
-"""
-GLM-4.6V (glm4v_moe) bridge for megatron.bridge.
-
-Registers `Glm4vMoeForConditionalGeneration` so that `AutoBridge.from_hf_pretrained`
-recognises GLM-4.6V checkpoints and can provide a Megatron-compatible VL model +
-weight mappings.
-
-Architecture:
-  HF vision encoder (Glm4vMoeVisionModel, replicated on first PP stage)
-  + Megatron GPTModel (MoE language model, standard M-RoPE)
-"""
-
-from __future__ import annotations
-
-import itertools
-import logging
-from copy import deepcopy
-from dataclasses import dataclass, field
-
-import torch
-from megatron.bridge.models.conversion.mapping_registry import MegatronMappingRegistry
-from megatron.bridge.models.conversion.model_bridge import MegatronModelBridge
-from megatron.bridge.models.conversion.param_mapping import AutoMapping, GatedMLPMapping, QKVMapping, ReplicatedMapping
-from megatron.bridge.models.qwen.qwen_provider import Qwen3MoEModelProvider
-from megatron.bridge.utils.common_utils import hook_hf_module_setattr_for_tp_grad_sync
-from megatron.core import parallel_state, tensor_parallel
-from megatron.core.models.gpt import GPTModel as MCoreGPTModel
-from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec
-from megatron.core.packed_seq_params import PackedSeqParams
-from megatron.core.transformer.module import MegatronModule
-
-logger = logging.getLogger(__name__)
-
-
-# ---------------------------------------------------------------------------
-# THD ↔ BSHD helpers (cf. Qwen3VL bridge)
-# ---------------------------------------------------------------------------
-def _thd_to_bshd(packed: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
-    """Unpack THD-format [1, T, ...] to BSHD [bs, max_seq, ...] using cu_seqlens."""
-    seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
-    max_seq = seqlens.max().item()
-    bs = len(cu_seqlens) - 1
-    out = packed.new_zeros(bs, max_seq, *packed.shape[2:])
-    for i, sl in enumerate(seqlens):
-        out[i, :sl] = packed[0, cu_seqlens[i] : cu_seqlens[i] + sl]
-    return out
-
-
-def _bshd_to_thd(unpacked: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor:
-    """Pack BSHD [bs, max_seq, ...] back to THD [1, T, ...]."""
-    seqlens = cu_seqlens[1:] - cu_seqlens[:-1]
-    total = cu_seqlens[-1].item()
-    out = unpacked.new_zeros(1, total, *unpacked.shape[2:])
-    for i, sl in enumerate(seqlens):
-        out[0, cu_seqlens[i] : cu_seqlens[i] + sl] = unpacked[i, :sl]
-    return out
-
-
-def _gather_input_ids_from_cp(
-    input_ids: torch.Tensor,
-    cu_seqlens: torch.Tensor,
-) -> torch.Tensor:
-    """Reconstruct full (global) input_ids from zigzag CP chunks.
-
-    With zigzag CP, each CP rank r holds chunks [r] and [2*cp_size-1-r] for
-    each sequence.  This function all-gathers across CP ranks and reassembles
-    the original token order so that position-ID computation sees the full
-    sequence.
-
-    Args:
-        input_ids: Local input_ids in THD format [1, T_local].
-        cu_seqlens: **Global** cumulative sequence lengths.
-
-    Returns:
-        Full input_ids in THD format [1, T_global].
-    """
-    cp_size = parallel_state.get_context_parallel_world_size()
-    if cp_size <= 1:
-        return input_ids
-
-    # all-gather local input_ids across CP ranks → list of [1, T_local] per rank
-    gathered = torch.distributed.nn.all_gather(
-        input_ids, group=parallel_state.get_context_parallel_group()
-    )  # list of cp_size tensors, each [1, T_local]
-
-    local_cu_seqlens = cu_seqlens // cp_size
-    num_seqs = len(cu_seqlens) - 1
-    whole_list = []
-    for i in range(num_seqs):
-        seqlen = (cu_seqlens[i + 1] - cu_seqlens[i]).item()
-        chunk_size = seqlen // 2 // cp_size
-        # First half: rank 0 chunk, rank 1 chunk, ..., rank cp_size-1 chunk
-        whole_list.extend(
-            gathered[cp_rank][0, local_cu_seqlens[i] : local_cu_seqlens[i] + chunk_size] for cp_rank in range(cp_size)
-        )
-        # Second half: rank cp_size-1 chunk, ..., rank 0 chunk (reversed)
-        whole_list.extend(
-            [
-                gathered[cp_rank][0, local_cu_seqlens[i] + chunk_size : local_cu_seqlens[i + 1]]
-                for cp_rank in range(cp_size)
-            ][::-1]
-        )
-    return torch.cat(whole_list).unsqueeze(0)  # [1, T_global]
-
-
-def _select_local_image_embeds(
-    full_input_ids: torch.Tensor,
-    cu_seqlens: torch.Tensor,
-    image_token_id: int,
-    image_embeds: torch.Tensor,
-    cp_rank: int,
-    cp_size: int,
-) -> torch.Tensor:
-    """Select the subset of *image_embeds* that falls in this CP rank's chunk.
-
-    With zigzag CP, each rank holds specific chunks of each packed sequence.
-    The vision encoder produces embeddings for ALL image tokens (ordered by
-    position in the full sequence).  This function returns only the embeddings
-    whose positions land in the local chunk.
-
-    Args:
-        full_input_ids: Reconstructed full input_ids [1, T_global].
-        cu_seqlens: **Global** cumulative sequence lengths.
-        image_token_id: Token id for image placeholder tokens.
-        image_embeds: Vision embeddings [N_total, hidden] for all image tokens.
-        cp_rank: This rank's position in the CP group.
-        cp_size: Total number of CP ranks.
-
-    Returns:
-        Subset of image_embeds [N_local, hidden] for this rank's image tokens.
-    """
-    device = full_input_ids.device
-    full_flat = full_input_ids[0]  # [T_global]
-    full_mask = full_flat == image_token_id
-
-    # Build boolean mask over T_global marking this rank's positions
-    T_global = full_flat.shape[0]
-    rank_mask = torch.zeros(T_global, dtype=torch.bool, device=device)
-
-    num_seqs = len(cu_seqlens) - 1
-    for i in range(num_seqs):
-        seq_start = cu_seqlens[i].item()
-        seqlen = (cu_seqlens[i + 1] - cu_seqlens[i]).item()
-        chunk_size = seqlen // (2 * cp_size)
-
-        # First-half chunk for this rank
-        first_start = seq_start + cp_rank * chunk_size
-        rank_mask[first_start : first_start + chunk_size] = True
-
-        # Second-half chunk for this rank (reversed order)
-        second_end = seq_start + seqlen - cp_rank * chunk_size
-        rank_mask[second_end - chunk_size : second_end] = True
-
-    # Image tokens that belong to this rank
-    local_image_mask = full_mask & rank_mask
-    n_local = local_image_mask.sum().item()
-
-    if n_local == 0:
-        return image_embeds[:0]  # empty slice preserving hidden dim
-    if n_local == image_embeds.shape[0]:
-        return image_embeds  # all image tokens are on this rank
-
-    # Map positions to indices in image_embeds via cumulative sum
-    image_cumsum = full_mask.long().cumsum(0)  # 1-indexed
-    local_positions = local_image_mask.nonzero(as_tuple=True)[0]
-    embed_indices = image_cumsum[local_positions] - 1
-    return image_embeds[embed_indices]
-
-
-# ---------------------------------------------------------------------------
-# Megatron VL Model
-# ---------------------------------------------------------------------------
-class Glm4vMoeVLModel(MegatronModule):
-    """GLM-4.6V vision-language model for Megatron training.
-
-    Wraps an HF vision encoder (only on first PP stage) together with a
-    standard Megatron Core GPTModel configured for M-RoPE.
-    """
-
-    def __init__(
-        self,
-        language_transformer_config,
-        language_transformer_layer_spec,
-        hf_vision_config,
-        parallel_output: bool = True,
-        pre_process: bool = True,
-        post_process: bool = True,
-    ) -> None:
-        super().__init__(config=language_transformer_config)
-
-        self.pre_process = pre_process
-        self.post_process = post_process
-        self.image_token_id = language_transformer_config.image_token_id
-        self.video_token_id = language_transformer_config.video_token_id
-        self.spatial_merge_size = language_transformer_config.spatial_merge_size
-
-        self.share_embeddings_and_output_weights = False
-
-        # Vision encoder -- only on the first pipeline stage
-        self.vision_model = None
-        if self.pre_process:
-            from transformers.models.glm4v_moe.modeling_glm4v_moe import Glm4vMoeVisionModel
-
-            self.vision_model = Glm4vMoeVisionModel._from_config(hf_vision_config)
-            # Freeze vision encoder — not trained during RL
-            self.vision_model.requires_grad_(False)
-            self.vision_model.eval()
-            hook_hf_module_setattr_for_tp_grad_sync(self.vision_model)
-            if torch.cuda.is_available():
-                self.vision_model = self.vision_model.to("cuda")
-
-        # Language model -- standard Megatron GPT with M-RoPE
-        self.language_model = MCoreGPTModel(
-            config=language_transformer_config,
-            transformer_layer_spec=language_transformer_layer_spec,
-            vocab_size=language_transformer_config.vocab_size,
-            max_sequence_length=language_transformer_config.language_max_sequence_length,
-            parallel_output=parallel_output,
-            position_embedding_type="mrope",
-            rotary_percent=language_transformer_config.rotary_percent,
-            pre_process=self.pre_process,
-            post_process=self.post_process,
-            rotary_base=language_transformer_config.rotary_base,
-            fp16_lm_cross_entropy=language_transformer_config.fp16_lm_cross_entropy,
-            share_embeddings_and_output_weights=language_transformer_config.share_embeddings_and_output_weights,
-            scatter_embedding_sequence_parallel=False,
-        )
-
-        self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights
-
-    # -- helpers required by Megatron pipeline engine -----------------------
-
-    def shared_embedding_or_output_weight(self):
-        return self.language_model.shared_embedding_or_output_weight()
-
-    def set_input_tensor(self, input_tensor):
-        if not isinstance(input_tensor, list):
-            input_tensor = [input_tensor]
-        assert len(input_tensor) == 1
-        if self.pre_process:
-            self.encoder_hidden_state = input_tensor[0]
-        else:
-            self.language_model.set_input_tensor(input_tensor[0])
-
-    # -- vision helpers -----------------------------------------------------
-
-    def _get_image_features(self, pixel_values, image_grid_thw):
-        """Run HF vision encoder and return flat image embeddings."""
-        pixel_values = pixel_values.to(dtype=self.vision_model.dtype)
-        with torch.no_grad():
-            return self.vision_model(pixel_values, grid_thw=image_grid_thw)
-
-    # -- M-RoPE position IDs -----------------------------------------------
-
-    @staticmethod
-    def _get_vision_position_ids(
-        start_position: int,
-        grid_thw,
-        temp_merge_size: int,
-        spatial_merge_size: int,
-        device,
-    ) -> torch.Tensor:
-        """Compute 3D positions for one image/video (ported from HF)."""
-        llm_grid_t = grid_thw[0].item() // temp_merge_size
-        llm_grid_h = grid_thw[1].item() // spatial_merge_size
-        llm_grid_w = grid_thw[2].item() // spatial_merge_size
-        n_tokens = llm_grid_h * llm_grid_w * llm_grid_t
-
-        pos_w = torch.arange(start_position, start_position + llm_grid_w, device=device)
-        pos_w = pos_w.repeat(llm_grid_h * llm_grid_t)
-        pos_h = torch.arange(start_position, start_position + llm_grid_h, device=device)
-        pos_h = pos_h.repeat_interleave(llm_grid_w * llm_grid_t)
-        pos_t = torch.full((n_tokens,), start_position, device=device, dtype=torch.long)
-        return torch.stack([pos_t, pos_h, pos_w], dim=0)  # [3, n_tokens]
-
-    def _compute_mrope_position_ids(
-        self,
-        input_ids_bshd: torch.Tensor,
-        image_grid_thw: torch.Tensor | None,
-    ) -> torch.Tensor:
-        """Compute 3D M-RoPE position IDs from input_ids in [bs, seq] format.
-
-        Image regions are detected by looking for consecutive runs of
-        ``image_token_id`` in each sequence — no ``mm_token_type_ids`` needed.
-        """
-        bs, seq_len = input_ids_bshd.shape
-        device = input_ids_bshd.device
-        spatial_merge_size = self.spatial_merge_size
-
-        position_ids = torch.zeros(3, bs, seq_len, dtype=torch.long, device=device)
-
-        if image_grid_thw is None or image_grid_thw.numel() == 0:
-            # Text-only: standard 1D positions replicated across 3 dims
-            pos = torch.arange(seq_len, device=device).unsqueeze(0).expand(bs, -1)
-            position_ids[0] = pos
-            position_ids[1] = pos
-            position_ids[2] = pos
-            return position_ids
-
-        grid_iter = iter(image_grid_thw)
-
-        for b in range(bs):
-            ids = input_ids_bshd[b]
-            is_image = ids == self.image_token_id
-
-            # Find contiguous groups: text (0) vs image (1)
-            token_types = is_image.long()
-            groups = []
-            for key, group in itertools.groupby(enumerate(token_types.tolist()), lambda x: x[1]):
-                g = list(group)
-                groups.append((key, g[0][0], g[-1][0] + 1))
-
-            current_pos = 0
-            pos_list = []
-            for modality, start, end in groups:
-                if modality == 0:
-                    # Text tokens
-                    n = end - start
-                    pos_list.append(torch.arange(n, device=device).view(1, -1).expand(3, -1) + current_pos)
-                    current_pos += n
-                else:
-                    # Image tokens
-                    grid_thw = next(grid_iter)
-                    temp_merge_size = grid_thw[0]
-                    vis_pos = self._get_vision_position_ids(
-                        current_pos,
-                        grid_thw,
-                        temp_merge_size,
-                        spatial_merge_size,
-                        device,
-                    )
-                    pos_list.append(vis_pos)
-                    current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
-
-            all_pos = torch.cat(pos_list, dim=1)  # [3, seq_for_this_sample]
-            position_ids[:, b, : all_pos.shape[1]] = all_pos
-
-        return position_ids
-
-    # -- forward ------------------------------------------------------------
-
-    def forward(
-        self,
-        input_ids: torch.Tensor,
-        position_ids: torch.Tensor = None,
-        attention_mask: torch.Tensor = None,
-        labels: torch.Tensor = None,
-        loss_mask: torch.Tensor = None,
-        inference_params=None,
-        packed_seq_params: PackedSeqParams = None,
-        extra_block_kwargs: dict = None,
-        # multimodal kwargs (unpacked from multimodal_train_inputs)
-        pixel_values: torch.Tensor = None,
-        image_grid_thw: torch.Tensor = None,
-        # unused VL kwargs that may come through
-        pixel_values_videos: torch.Tensor = None,
-        video_grid_thw: torch.Tensor = None,
-        mm_token_type_ids: torch.Tensor = None,
-        **kwargs,
-    ) -> torch.Tensor:
-        assert pixel_values_videos is None, "Video not supported yet"
-        assert inference_params is None, "Inference not supported"
-
-        # -- Extract cu_seqlens and CP info early (needed for both vision scatter and M-RoPE) --
-        cu_seqlens = None
-        if packed_seq_params is not None:
-            cu_seqlens = (
-                packed_seq_params.cu_seqlens_q_padded
-                if packed_seq_params.cu_seqlens_q_padded is not None
-                else packed_seq_params.cu_seqlens_q
-            )
-        cp_size = parallel_state.get_context_parallel_world_size()
-        full_input_ids = None  # cached for reuse between vision scatter and M-RoPE
-
-        combined_embeddings = None
-
-        if self.pre_process:
-            # 1. Text embeddings from language model embedding layer
-            combined_embeddings = self.language_model.embedding(
-                input_ids=input_ids,
-                position_ids=None,
-            ).clone()  # [seq, batch, hidden]
-
-            # 2. Vision encoding + masked scatter
-            if pixel_values is not None and image_grid_thw is not None:
-                image_embeds = self._get_image_features(pixel_values, image_grid_thw)
-                image_embeds = image_embeds.to(combined_embeddings.device, combined_embeddings.dtype)
-
-                # With CP > 1, input_ids is a local chunk but pixel_values
-                # cover ALL images.  Select only the embeddings whose tokens
-                # land in this rank's zigzag portion.
-                if cp_size > 1 and cu_seqlens is not None:
-                    full_input_ids = _gather_input_ids_from_cp(input_ids, cu_seqlens)
-                    cp_rank = parallel_state.get_context_parallel_rank()
-                    image_embeds = _select_local_image_embeds(
-                        full_input_ids,
-                        cu_seqlens,
-                        self.image_token_id,
-                        image_embeds,
-                        cp_rank,
-                        cp_size,
-                    )
-
-                image_mask = (input_ids == self.image_token_id).contiguous()
-                # Scatter: [seq, bs, hidden] → [bs, seq, hidden]
-                combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()
-                if image_mask.any():
-                    combined_embeddings[image_mask] = image_embeds
-                combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()
-
-            # Scatter to sequence-parallel region if needed
-            if self.config.sequence_parallel:
-                combined_embeddings = tensor_parallel.scatter_to_sequence_parallel_region(combined_embeddings)
-                combined_embeddings = combined_embeddings.contiguous()
-
-        # 3. Compute M-RoPE position IDs
-        # position_ids must be available on ALL PP stages for rotary embeddings.
-        # On stage 0, compute from input_ids. Then broadcast to other stages.
-        pp_size = parallel_state.get_pipeline_model_parallel_world_size()
-
-        if position_ids is None:
-            if self.pre_process:
-                # First PP stage: compute position_ids from input_ids.
-                # With CP > 1, input_ids is a local chunk; reconstruct full
-                # sequence so that _compute_mrope_position_ids sees all tokens
-                # (image token positions affect the M-RoPE IDs).
-                if cu_seqlens is not None:
-                    if cp_size > 1:
-                        if full_input_ids is None:
-                            full_input_ids = _gather_input_ids_from_cp(input_ids, cu_seqlens)
-                    else:
-                        full_input_ids = input_ids
-                    input_ids_bshd = _thd_to_bshd(full_input_ids, cu_seqlens)
-                    pos_bshd = self._compute_mrope_position_ids(input_ids_bshd, image_grid_thw)
-                    pos_packed = _bshd_to_thd(pos_bshd.permute(1, 2, 0), cu_seqlens)
-                    position_ids = pos_packed.permute(2, 0, 1).contiguous()  # [3, 1, T_global]
-                else:
-                    position_ids = self._compute_mrope_position_ids(input_ids, image_grid_thw)
-            else:
-                # Non-first PP stage: allocate buffer with correct shape
-                if cu_seqlens is not None:
-                    T = cu_seqlens[-1].item()
-                    position_ids = torch.zeros(3, 1, T, dtype=torch.long, device=torch.cuda.current_device())
-                else:
-                    raise NotImplementedError(
-                        "Non-THD position_ids broadcast not yet supported for non-first PP stages"
-                    )
-
-            # Broadcast position_ids from first to all PP stages
-            if pp_size > 1:
-                src = parallel_state.get_pipeline_model_parallel_first_rank()
-                torch.distributed.broadcast(
-                    position_ids,
-                    src=src,
-                    group=parallel_state.get_pipeline_model_parallel_group(),
-                )
-
-        # 4. Language model forward (pass decoder_input to skip re-embedding)
-        output = self.language_model(
-            input_ids=None,
-            position_ids=position_ids,
-            attention_mask=attention_mask,
-            decoder_input=combined_embeddings,
-            labels=labels,
-            loss_mask=loss_mask,
-            inference_params=inference_params,
-            packed_seq_params=packed_seq_params,
-            **(extra_block_kwargs or {}),
-        )
-
-        return output
-
-
-# ---------------------------------------------------------------------------
-# Model Provider (dataclass that doubles as TransformerConfig)
-# ---------------------------------------------------------------------------
-@dataclass
-class Glm4vMoeVLModelProvider(Qwen3MoEModelProvider):
-    """Provider that creates Glm4vMoeVLModel.
-
-    Inherits from Qwen3MoEModelProvider to reuse MoE + TransformerConfig infra.
-    Defined at module level (not inside a function) so that the class is
-    picklable -- megatron-bridge broadcasts config objects across PP ranks
-    via ``torch.distributed.broadcast_object_list`` which requires pickling.
-    """
-
-    # GLM-4.6V specific config
-    image_token_id: int = 151363
-    video_token_id: int = 151364
-    spatial_merge_size: int = 2
-
-    # Vision config (stored as HF config object)
-    hf_vision_config: object = None
-    hf_text_config: object = None
-
-    # M-RoPE
-    position_embedding_type: str = "mrope"
-    mrope_section: list[int] = field(default_factory=lambda: [8, 12, 12])
-    scatter_embedding_sequence_parallel: bool = False
-
-    # Language model sequence length
-    language_max_sequence_length: int = 131072
-
-    def provide(self, pre_process=None, post_process=None, vp_stage=None):
-        """Create a Glm4vMoeVLModel instance."""
-
-        # Resolve PP stage flags
-        if pre_process is None:
-            pre_process = parallel_state.is_pipeline_first_stage(ignore_virtual=False, vp_stage=vp_stage)
-        if post_process is None:
-            post_process = parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=vp_stage)
-
-        # Build per-layer specs respecting moe_layer_freq (layer 0 = dense, rest = MoE)
-        transformer_layer_spec = get_gpt_decoder_block_spec(
-            config=self,
-            use_transformer_engine=True,
-            vp_stage=vp_stage,
-        )
-
-        model = Glm4vMoeVLModel(
-            language_transformer_config=self,
-            language_transformer_layer_spec=transformer_layer_spec,
-            hf_vision_config=self.hf_vision_config,
-            parallel_output=True,
-            pre_process=pre_process,
-            post_process=post_process,
-        )
-
-        return model
-
-
-# ---------------------------------------------------------------------------
-# Bridge
-# ---------------------------------------------------------------------------
-try:
-    from transformers import Glm4vMoeForConditionalGeneration as _Glm4vMoeHF
-except ImportError:
-    _Glm4vMoeHF = "Glm4vMoeForConditionalGeneration"
-
-
-@MegatronModelBridge.register_bridge(source=_Glm4vMoeHF, target=Glm4vMoeVLModel)
-class Glm4vMoeBridge(MegatronModelBridge):
-    """Bridge between HuggingFace GLM-4.6V and the Megatron VL model."""
-
-    def provider_bridge(self, hf_pretrained):
-        """Create a Glm4vMoeVLModelProvider from HF config."""
-        hf_config = hf_pretrained.config
-        text_config = hf_config.text_config
-        vision_config = deepcopy(hf_config.vision_config)
-
-        model_dtype = self.dtype_from_hf(text_config, default=torch.bfloat16)
-        vision_config.torch_dtype = model_dtype
-
-        ProviderClass = Glm4vMoeVLModelProvider
-
-        rope_params = getattr(text_config, "rope_parameters", {}) or {}
-        mrope_section = rope_params.get("mrope_section", [8, 12, 12])
-        rotary_base = rope_params.get("rope_theta", 500000)
-        partial_rotary_factor = rope_params.get("partial_rotary_factor", 0.5)
-
-        # Determine MoE layer frequency
-        first_k_dense = getattr(text_config, "first_k_dense_replace", 1)
-        num_layers = text_config.num_hidden_layers
-        # Build moe_layer_freq list: first_k_dense dense layers + rest MoE
-        moe_layer_freq_list = [0] * first_k_dense + [1] * (num_layers - first_k_dense)
-
-        # Shared expert intermediate size
-        n_shared = getattr(text_config, "n_shared_experts", 1)
-        moe_ffn = getattr(text_config, "moe_intermediate_size", 1408)
-        shared_expert_intermediate = moe_ffn * n_shared
-
-        provider = ProviderClass(
-            # Language model configuration
-            num_layers=num_layers,
-            hidden_size=text_config.hidden_size,
-            ffn_hidden_size=text_config.intermediate_size,
-            num_attention_heads=text_config.num_attention_heads,
-            num_query_groups=text_config.num_key_value_heads,
-            kv_channels=getattr(text_config, "head_dim", 128),
-            init_method_std=text_config.initializer_range,
-            layernorm_epsilon=text_config.rms_norm_eps,
-            gated_linear_unit=True,
-            make_vocab_size_divisible_by=self.make_vocab_size_divisible_by(text_config.vocab_size),
-            rotary_base=rotary_base,
-            rotary_percent=partial_rotary_factor,
-            share_embeddings_and_output_weights=getattr(text_config, "tie_word_embeddings", False),
-            vocab_size=text_config.vocab_size,
-            seq_length=text_config.max_position_embeddings,
-            fp16=(model_dtype == torch.float16),
-            bf16=(model_dtype == torch.bfloat16),
-            params_dtype=model_dtype,
-            # MoE configuration
-            num_moe_experts=getattr(text_config, "n_routed_experts", 128),
-            moe_router_topk=getattr(text_config, "num_experts_per_tok", 8),
-            moe_ffn_hidden_size=moe_ffn,
-            moe_shared_expert_intermediate_size=shared_expert_intermediate,
-            moe_layer_freq=moe_layer_freq_list,
-            moe_grouped_gemm=True,
-            moe_router_load_balancing_type="seq_aux_loss",
-            moe_aux_loss_coeff=0,
-            moe_router_score_function="sigmoid",
-            moe_router_pre_softmax=True,
-            moe_router_enable_expert_bias=True,
-            moe_router_dtype="fp32",
-            # Attention
-            add_qkv_bias=getattr(text_config, "attention_bias", True),
-            qk_layernorm=getattr(text_config, "qk_layernorm", False) or getattr(text_config, "use_qk_norm", False),
-            # M-RoPE
-            mrope_section=mrope_section,
-            position_embedding_type="mrope",
-            scatter_embedding_sequence_parallel=False,
-            # Vision
-            hf_vision_config=vision_config,
-            hf_text_config=text_config,
-            image_token_id=getattr(hf_config, "image_token_id", 151363),
-            video_token_id=getattr(hf_config, "video_token_id", 151364),
-            spatial_merge_size=getattr(hf_config.vision_config, "spatial_merge_size", 2),
-            language_max_sequence_length=text_config.max_position_embeddings,
-        )
-
-        return provider
-
-    def mapping_registry(self) -> MegatronMappingRegistry:
-        """Weight mappings from HF GLM-4.6V to Megatron format.
-
-        Follows GLM-4.5 bridge pattern with language_model prefix for VL model.
-        Layer 0 is dense, layers 1-45 are MoE. The mapping framework handles
-        missing keys gracefully (warnings for non-existent params).
-        """
-        param_mappings = {
-            # Embeddings and output
-            "language_model.embedding.word_embeddings.weight": "model.language_model.embed_tokens.weight",
-            "language_model.output_layer.weight": "lm_head.weight",
-            "language_model.decoder.final_layernorm.weight": "model.language_model.norm.weight",
-            # Attention: input layernorm (fused with TE)
-            "language_model.decoder.layers.*.self_attention.linear_qkv.layer_norm_weight": "model.language_model.layers.*.input_layernorm.weight",
-            # Attention: separate input layernorm (quantization layer spec)
-            "language_model.decoder.layers.*.input_layernorm.weight": "model.language_model.layers.*.input_layernorm.weight",
-            # Attention output
-            "language_model.decoder.layers.*.self_attention.linear_proj.weight": "model.language_model.layers.*.self_attn.o_proj.weight",
-            # Post-attention layernorm:
-            #   MoE layers → pre_mlp_layernorm, Dense layer → linear_fc1.layer_norm_weight (fused)
-            "language_model.decoder.layers.*.pre_mlp_layernorm.weight": "model.language_model.layers.*.post_attention_layernorm.weight",
-            "language_model.decoder.layers.*.mlp.linear_fc1.layer_norm_weight": "model.language_model.layers.*.post_attention_layernorm.weight",
-            # Dense MLP output (layer 0)
-            "language_model.decoder.layers.*.mlp.linear_fc2.weight": "model.language_model.layers.*.mlp.down_proj.weight",
-            # MoE router
-            "language_model.decoder.layers.*.mlp.router.weight": "model.language_model.layers.*.mlp.gate.weight",
-            "language_model.decoder.layers.*.mlp.router.expert_bias": "model.language_model.layers.*.mlp.gate.e_score_correction_bias",
-            # MoE expert output (TEGroupedMLP format: weight* suffix)
-            "language_model.decoder.layers.*.mlp.experts.linear_fc2.weight*": "model.language_model.layers.*.mlp.experts.*.down_proj.weight",
-            # MoE shared expert output
-            "language_model.decoder.layers.*.mlp.shared_experts.linear_fc2.weight": "model.language_model.layers.*.mlp.shared_experts.down_proj.weight",
-        }
-
-        mapping_list = []
-        for megatron_param, hf_param in param_mappings.items():
-            mapping_list.append(AutoMapping(megatron_param=megatron_param, hf_param=hf_param))
-
-        mapping_list.extend(
-            [
-                # Vision model weights — replicated directly
-                ReplicatedMapping(
-                    megatron_param="vision_model.**",
-                    hf_param="model.visual.**",
-                ),
-                # QKV weight and bias
-                QKVMapping(
-                    megatron_param="language_model.decoder.layers.*.self_attention.linear_qkv.weight",
-                    q="model.language_model.layers.*.self_attn.q_proj.weight",
-                    k="model.language_model.layers.*.self_attn.k_proj.weight",
-                    v="model.language_model.layers.*.self_attn.v_proj.weight",
-                ),
-                QKVMapping(
-                    megatron_param="language_model.decoder.layers.*.self_attention.linear_qkv.bias",
-                    q="model.language_model.layers.*.self_attn.q_proj.bias",
-                    k="model.language_model.layers.*.self_attn.k_proj.bias",
-                    v="model.language_model.layers.*.self_attn.v_proj.bias",
-                ),
-                # Dense MLP gate+up (layer 0)
-                GatedMLPMapping(
-                    megatron_param="language_model.decoder.layers.*.mlp.linear_fc1.weight",
-                    gate="model.language_model.layers.*.mlp.gate_proj.weight",
-                    up="model.language_model.layers.*.mlp.up_proj.weight",
-                ),
-                # MoE expert gate+up (TEGroupedMLP format)
-                GatedMLPMapping(
-                    megatron_param="language_model.decoder.layers.*.mlp.experts.linear_fc1.weight*",
-                    gate="model.language_model.layers.*.mlp.experts.*.gate_proj.weight",
-                    up="model.language_model.layers.*.mlp.experts.*.up_proj.weight",
-                ),
-                # MoE expert gate+up (SequentialMLP format, for quantization)
-                GatedMLPMapping(
-                    megatron_param="language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc1.weight",
-                    gate="model.language_model.layers.*.mlp.experts.*.gate_proj.weight",
-                    up="model.language_model.layers.*.mlp.experts.*.up_proj.weight",
-                ),
-                AutoMapping(
-                    megatron_param="language_model.decoder.layers.*.mlp.experts.local_experts.*.linear_fc2.weight",
-                    hf_param="model.language_model.layers.*.mlp.experts.*.down_proj.weight",
-                ),
-                # MoE shared expert gate+up
-                GatedMLPMapping(
-                    megatron_param="language_model.decoder.layers.*.mlp.shared_experts.linear_fc1.weight",
-                    gate="model.language_model.layers.*.mlp.shared_experts.gate_proj.weight",
-                    up="model.language_model.layers.*.mlp.shared_experts.up_proj.weight",
-                ),
-            ]
-        )
-
-        return MegatronMappingRegistry(*mapping_list)
diff --git a/slime_plugins/models/hf_attention.py b/slime_plugins/models/hf_attention.py
index d5cf619724..c353ae7b29 100644
--- a/slime_plugins/models/hf_attention.py
+++ b/slime_plugins/models/hf_attention.py
@@ -1,5 +1,3 @@
-import json
-import os
 from abc import ABC, abstractmethod
 
 import torch
@@ -8,56 +6,7 @@
 from megatron.core.inference.contexts import BaseInferenceContext
 from megatron.core.packed_seq_params import PackedSeqParams
 from megatron.core.transformer.module import MegatronModule
-
-
-def _load_hf_config(checkpoint_path):
-    """Load HF config with fallback for unsupported model types."""
-    try:
-        from transformers import AutoConfig
-
-        return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True)
-    except (ValueError, KeyError):
-        config_path = os.path.join(checkpoint_path, "config.json")
-        with open(config_path) as f:
-            config_dict = json.load(f)
-
-        _DTYPE_MAP = {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}
-
-        def _fix_dtype(d):
-            if "torch_dtype" in d:
-                d["torch_dtype"] = _DTYPE_MAP.get(d["torch_dtype"], d["torch_dtype"])
-            if "dtype" in d:
-                d["dtype"] = _DTYPE_MAP.get(d["dtype"], d["dtype"])
-
-        _fix_dtype(config_dict)
-        ns = type("HFConfig", (), config_dict)()
-        if "text_config" in config_dict:
-            _fix_dtype(config_dict["text_config"])
-            ns.text_config = type("TextConfig", (), config_dict["text_config"])()
-        return ns
-
-
-class _AllGatherForDuplicatedComputation(torch.autograd.Function):
-    """All-gather whose backward just returns the local gradient slice (no reduce).
-
-    Use this instead of ``dist.nn.all_gather`` when the computation after the
-    gather is *duplicated* across ranks (same weights, same full input →
-    identical gradients).  The default ``all_gather`` backward performs a
-    reduce-scatter, which incorrectly sums ``world_size`` identical copies of
-    the gradient.
-    """
-
-    @staticmethod
-    def forward(ctx, x, group):
-        ctx.group = group
-        ctx.rank = dist.get_rank(group=group)
-        out = [torch.empty_like(x) for _ in range(dist.get_world_size(group=group))]
-        dist.all_gather(out, x.contiguous(), group=group)
-        return tuple(out)
-
-    @staticmethod
-    def backward(ctx, *grads):
-        return grads[ctx.rank], None
+from transformers import AutoConfig
 
 
 class HuggingfaceAttention(MegatronModule, ABC):
@@ -81,7 +30,7 @@ def __init__(
         # Note that megatron layer_number starts at 1
         self.layer_number = layer_number
         self.hf_layer_idx = layer_number - 1
-        self.hf_config = _load_hf_config(args.hf_checkpoint)
+        self.hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
         # hardcode to fa2 at the moment.
         self.hf_config._attn_implementation = "flash_attention_2"
 
@@ -105,22 +54,15 @@ def forward(
         cu_seqlens = packed_seq_params.cu_seqlens_q
 
         if self.args.sequence_parallel:
-            # tensor_parallel_output_grad=False: the linear attention after this
-            # gather is NOT TP-sharded (duplicated on all ranks), so the backward
-            # should split (not reduce-scatter) to avoid inflating gradients by TP.
             hidden_states = tensor_parallel.gather_from_sequence_parallel_region(
-                hidden_states,
-                tensor_parallel_output_grad=False,
-                group=mpu.get_tensor_model_parallel_group(),
+                hidden_states, group=mpu.get_tensor_model_parallel_group()
             )
 
         if mpu.get_context_parallel_world_size() > 1:
             cp_size = mpu.get_context_parallel_world_size()
-            # Use custom all-gather whose backward returns local gradient
-            # instead of reduce-scatter, since the computation is duplicated.
-            hidden_states_list = _AllGatherForDuplicatedComputation.apply(
+            hidden_states_list = dist.nn.all_gather(
                 hidden_states,
-                mpu.get_context_parallel_group(),
+                group=mpu.get_context_parallel_group(),
             )
 
             # TODO: preprocess this for each batch to prevent tolist in the training step
diff --git a/slime_plugins/models/qwen3_5.py b/slime_plugins/models/qwen3_5.py
index eb8cff0a31..33ff1a45b5 100644
--- a/slime_plugins/models/qwen3_5.py
+++ b/slime_plugins/models/qwen3_5.py
@@ -1,4 +1,6 @@
 import copy
+import json
+import os
 
 import torch
 import torch.nn as nn
@@ -9,13 +11,32 @@
 from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
 from transformers.activations import ACT2FN
 
+
+def _load_hf_config(checkpoint_path):
+    """Load HF config, handling cases where transformers doesn't know the model type."""
+    try:
+        from transformers import AutoConfig
+
+        return AutoConfig.from_pretrained(checkpoint_path, trust_remote_code=True)
+    except (ValueError, KeyError):
+        # Fallback: load config.json directly as a SimpleNamespace
+        config_path = os.path.join(checkpoint_path, "config.json")
+        with open(config_path) as f:
+            config_dict = json.load(f)
+        # If there's a text_config, also make it a namespace
+        ns = type("HFConfig", (), config_dict)()
+        if "text_config" in config_dict:
+            ns.text_config = type("TextConfig", (), config_dict["text_config"])()
+        return ns
+
+
 try:
     from fla.modules import FusedRMSNormGated, ShortConvolution
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule
 except ImportError:
     pass
 
-from .hf_attention import HuggingfaceAttention, _load_hf_config
+from .hf_attention import HuggingfaceAttention
 
 
 def _get_text_config(hf_config):
@@ -76,7 +97,7 @@ def __init__(self, config, layer_idx: int):
             eps=self.layer_norm_epsilon,
             activation=self.activation,
             device=torch.cuda.current_device(),
-            dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
+            dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(),
         )
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
@@ -127,7 +148,6 @@ def forward(
             initial_state=None,
             output_final_state=False,
             use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
         )
 
         z_shape_og = z.shape
@@ -205,14 +225,6 @@ def get_qwen3_5_spec(args, config, vp_stage):
     hf_config = _load_hf_config(args.hf_checkpoint)
     text_config = _get_text_config(hf_config)
 
-    # Compute layer_types if the config class doesn't expose it
-    if not hasattr(text_config, "layer_types"):
-        interval = getattr(text_config, "full_attention_interval", 4)
-        n = text_config.num_hidden_layers
-        text_config.layer_types = [
-            "full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(n)
-        ]
-
     for layer_id in range(num_layers_to_build):
         if text_config.layer_types[layer_id + offset] == "linear_attention":
             layer_specs = copy.deepcopy(transformer_layer_spec.layer_specs[layer_id])
diff --git a/slime_plugins/models/qwen3_next.py b/slime_plugins/models/qwen3_next.py
index 3060a945c6..0a42e4d572 100644
--- a/slime_plugins/models/qwen3_next.py
+++ b/slime_plugins/models/qwen3_next.py
@@ -7,10 +7,9 @@
 from megatron.core.transformer.spec_utils import ModuleSpec
 from megatron.core.transformer.transformer_block import get_num_layers_to_build
 from megatron.core.transformer.transformer_layer import get_transformer_layer_offset
+from transformers import AutoConfig
 from transformers.activations import ACT2FN
 
-from .hf_attention import _load_hf_config
-
 try:
     from fla.modules import FusedRMSNormGated, ShortConvolution
     from fla.ops.gated_delta_rule import chunk_gated_delta_rule
@@ -69,7 +68,7 @@ def __init__(self, config, layer_idx: int):
             eps=self.layer_norm_epsilon,
             activation=self.activation,
             device=torch.cuda.current_device(),
-            dtype=config.dtype if config.dtype is not None else torch.get_default_dtype(),
+            dtype=config.dtype if config.dtype is not None else torch.get_current_dtype(),
         )
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
@@ -149,7 +148,6 @@ def forward(
             initial_state=None,
             output_final_state=False,
             use_qk_l2norm_in_kernel=True,
-            cu_seqlens=cu_seqlens,
         )
 
         z_shape_og = z.shape
@@ -215,13 +213,7 @@ def get_qwen3_next_spec(args, config, vp_stage):
     num_layers_to_build = get_num_layers_to_build(config, vp_stage=vp_stage)
     offset = get_transformer_layer_offset(config, vp_stage=vp_stage)
 
-    hf_config = _load_hf_config(args.hf_checkpoint)
-
-    # Compute layer_types if the config class doesn't expose it
-    if not hasattr(hf_config, "layer_types"):
-        interval = getattr(hf_config, "full_attention_interval", 4)
-        n = hf_config.num_hidden_layers
-        hf_config.layer_types = ["full_attention" if (i + 1) % interval == 0 else "linear_attention" for i in range(n)]
+    hf_config = AutoConfig.from_pretrained(args.hf_checkpoint, trust_remote_code=True)
 
     for layer_id in range(num_layers_to_build):
         if hf_config.layer_types[layer_id + offset] == "linear_attention":
diff --git a/tests/plugin_contracts/test_plugin_generate_contracts.py b/tests/plugin_contracts/test_plugin_generate_contracts.py
index 9cf7af3a32..d63a69f45f 100644
--- a/tests/plugin_contracts/test_plugin_generate_contracts.py
+++ b/tests/plugin_contracts/test_plugin_generate_contracts.py
@@ -173,29 +173,5 @@ def test_custom_generate_function_path_supports_user_override(patch_generate_sta
     assert_sample_contract(result)
 
 
-def test_generate_and_rm_group_rm_accepts_list_result_from_custom_generate(patch_generate_state, monkeypatch):
-    sglang_rollout = patch_generate_state
-
-    async def custom_generate_list(args, sample: Sample, sampling_params: dict):
-        sample.status = Sample.Status.COMPLETED
-        sibling = Sample(index=1, prompt="prompt-1", status=Sample.Status.COMPLETED)
-        return [sample, sibling]
-
-    monkeypatch.setattr(sglang_rollout, "load_function", lambda _path: custom_generate_list)
-
-    result = asyncio.run(
-        generate_and_rm(
-            make_args(custom_generate_function_path="plugin_contracts.fake_generate", group_rm=True),
-            Sample(index=0, prompt="prompt-0"),
-            sampling_params={"temperature": 0.3},
-            evaluation=False,
-        )
-    )
-
-    assert isinstance(result, list)
-    assert len(result) == 2
-    assert all(isinstance(sample, Sample) for sample in result)
-
-
 if __name__ == "__main__":
     run_contract_test_file()
diff --git a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
index a8380feecc..9c74876d96 100644
--- a/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
+++ b/tests/plugin_contracts/test_plugin_runtime_hook_contracts.py
@@ -69,7 +69,7 @@ def reference_convert_samples_to_train_data(args, samples):
     }
 
 
-def reference_rollout_data_postprocess(args, rollout_id, rollout_data) -> None:
+def reference_rollout_data_postprocess(args) -> None:
     args.rollout_data_postprocess_called = True
 
 
@@ -124,7 +124,7 @@ def invoke_convert_samples_to_train_data(fn):
 
 def invoke_rollout_data_postprocess(fn):
     args = type("Args", (), {})()
-    assert fn(args, 0, {}) is None
+    assert fn(args) is None
     assert args.rollout_data_postprocess_called is True
 
 
@@ -170,8 +170,8 @@ def invoke_rollout_data_postprocess(fn):
         "ROLLOUT_DATA_POSTPROCESS_PATH",
         "plugin_contracts.test_plugin_runtime_hook_contracts.reference_rollout_data_postprocess",
         "slime/backends/megatron_utils/actor.py",
-        "self.rollout_data_postprocess(self.args, rollout_id, rollout_data)",
-        ("args", "rollout_id", "rollout_data"),
+        "self.rollout_data_postprocess(self.args)",
+        ("args",),
         invoke_rollout_data_postprocess,
     ),
 ]
diff --git a/tests/test_fsdp_import.py b/tests/test_fsdp_import.py
new file mode 100644
index 0000000000..66b6861ed1
--- /dev/null
+++ b/tests/test_fsdp_import.py
@@ -0,0 +1,9 @@
+import pytest
+
+
+def test_fsdp_import():
+    try:
+        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+    except ImportError:
+        pytest.skip("FSDP not available in this environment")
+    assert FSDP is not None
diff --git a/tests/test_fused_experts_backward.py b/tests/test_fused_experts_backward.py
new file mode 100644
index 0000000000..945a3a0480
--- /dev/null
+++ b/tests/test_fused_experts_backward.py
@@ -0,0 +1,593 @@
+"""
+Test script to compare Triton backward implementation with Python backward implementation.
+
+This test compares:
+1. Triton implementation (from fused_experts.py) - uses invoke_fused_moe_backward_kernel
+2. Python reference implementation (defined in this file) - uses pure PyTorch operations
+"""
+
+import pytest
+import torch
+
+# ============================================================================
+# Python Reference Implementation (Pure PyTorch)
+# ============================================================================
+
+
+class GateUpProjFunctionPython(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        hidden_states: torch.Tensor,
+        w1: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        num_tokens, D_in = hidden_states.shape
+        E, N, K = w1.shape
+        assert D_in == K, f"hidden_states dim {D_in} != w1 dim {K}"
+
+        topk = topk_ids.shape[1]
+
+        # Output: (num_tokens * topk, N)
+        intermediate_cache1 = torch.empty(
+            (num_tokens * topk, N),
+            device=hidden_states.device,
+            dtype=hidden_states.dtype,
+        )
+
+        # Python implementation: iterate over tokens and their topk experts
+        # For each token t and expert k:
+        #   intermediate_cache1[t*topk + k] = hidden_states[t] @ w1[expert_id].T
+        for t in range(num_tokens):
+            for k in range(topk):
+                expert_id = topk_ids[t, k].item()
+                x_t = hidden_states[t]  # shape: (D_in,)
+                W1_e = w1[expert_id]  # shape: (N, K)
+                intermediate_cache1[t * topk + k] = x_t @ W1_e.T
+
+        ctx.save_for_backward(hidden_states, w1, topk_weights, topk_ids)
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache1
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for GateUpProjFunction - Pure Python implementation.
+
+        Forward: output = input @ w1 (without topk_weight multiplication)
+        Backward:
+            - grad_hidden_states = grad_output @ w1
+            - grad_w1 = grad_output.T @ input (note: transposed)
+            - grad_topk_weights = zeros (not needed in this stage)
+
+        Args:
+            grad_output: shape (num_tokens * topk, N)
+
+        Returns:
+            (grad_hidden_states, grad_w1, grad_topk_weights, None)
+        """
+        hidden_states, w1, topk_weights, topk_ids = ctx.saved_tensors
+        topk = ctx.topk
+
+        num_tokens, D_in = hidden_states.shape
+        E, N, _ = w1.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_hidden_states = torch.zeros_like(hidden_states)
+        # Use float32 for grad_w1 accumulation to avoid bfloat16 precision loss
+        grad_w1 = torch.zeros(w1.shape, dtype=torch.float32, device=w1.device)
+        # GateUpProj stage doesn't compute topk_weights gradient
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_hidden_states = hidden_states[begin_chunk_idx:end_chunk_idx]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx * topk : end_chunk_idx * topk]
+
+            # 1. Calculate grad_hidden_states: grad_output @ w1
+            # For each token t and expert k:
+            #   grad_hidden_states[t] += grad_output[t*topk+k] @ w1[expert_id]
+            for t in range(curr_num_tokens):
+                for k in range(topk):
+                    expert_id = curr_topk_ids[t, k].item()
+                    grad_y_tk = curr_grad_output[t * topk + k]  # shape: (N,)
+                    W1_e = w1[expert_id]  # shape: (N, D_in)
+                    # grad_x: (N,) @ (N, D_in) -> (D_in,)
+                    grad_hidden_states[begin_chunk_idx + t] += grad_y_tk @ W1_e
+
+            # 2. Calculate grad_w1: input.T @ grad_output
+            # For each token t and expert k:
+            #   grad_w1[expert_id] += input[t].T @ grad_output[t*topk+k]
+            # Which is: grad_w1[expert_id] += outer(grad_output[t*topk+k], input[t])
+            for t in range(curr_num_tokens):
+                for k in range(topk):
+                    expert_id = curr_topk_ids[t, k].item()
+                    x_t = curr_hidden_states[t]  # shape: (D_in,)
+                    grad_y_tk = curr_grad_output[t * topk + k]  # shape: (N,)
+                    # grad_W1: outer(grad_y_tk, x_t) -> (N, D_in)
+                    # Accumulate in float32
+                    grad_w1[expert_id] += torch.outer(grad_y_tk, x_t).to(torch.float32)
+
+        # Convert grad_w1 back to original dtype (bfloat16)
+        grad_w1 = grad_w1.to(hidden_states.dtype)
+
+        return grad_hidden_states, grad_w1, grad_topk_weights, None
+
+
+class DownProjFunctionPython(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        intermediate_cache2: torch.Tensor,
+        w2: torch.Tensor,
+        topk_weights: torch.Tensor,
+        topk_ids: torch.Tensor,
+    ):
+        total_tokens, intermediate_size = intermediate_cache2.shape
+        topk = topk_ids.shape[1]
+        num_tokens = total_tokens // topk
+        E, hidden_size, K = w2.shape
+        assert intermediate_size == K, f"intermediate_cache2 dim {intermediate_size} != w2 dim {K}"
+
+        # Output: (num_tokens, topk, hidden_size)
+        intermediate_cache3 = torch.empty(
+            (num_tokens, topk, hidden_size),
+            device=intermediate_cache2.device,
+            dtype=intermediate_cache2.dtype,
+        )
+
+        # Python implementation: iterate over tokens and their topk experts
+        # For each token t and expert k:
+        #   intermediate_cache3[t, k] = topk_weights[t, k] * (intermediate_cache2[t*topk+k] @ w2[expert_id].T)
+        for t in range(num_tokens):
+            for k in range(topk):
+                expert_id = topk_ids[t, k].item()
+                x_tk = intermediate_cache2[t * topk + k]  # shape: (intermediate_size,)
+                W2_e = w2[expert_id]  # shape: (hidden_size, intermediate_size)
+                weight_tk = topk_weights[t, k]  # scalar
+
+                intermediate_cache3[t, k] = weight_tk * (x_tk @ W2_e.T)
+
+        ctx.save_for_backward(intermediate_cache2, w2, topk_weights, topk_ids)
+        ctx.num_tokens = num_tokens
+        ctx.topk = topk
+
+        return intermediate_cache3
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """
+        Backward pass for DownProjFunction - Pure Python implementation.
+
+        Forward: output = topk_weights * (input @ w2) (with topk_weight multiplication)
+        Backward:
+            - grad_intermediate_cache2 = topk_weights * (grad_output @ w2)
+            - grad_w2 = topk_weights * (grad_output.T @ intermediate_cache2)
+            - grad_topk_weights = dot(grad_output, forward_output_before_weighting)
+
+        Args:
+            grad_output: shape (num_tokens, topk, hidden_size)
+
+        Returns:
+            (grad_intermediate_cache2, grad_w2, grad_topk_weights, None)
+        """
+        intermediate_cache2, w2, topk_weights, topk_ids = ctx.saved_tensors
+        num_tokens = ctx.num_tokens
+        topk = ctx.topk
+
+        E, hidden_size, intermediate_size = w2.shape
+        CHUNK_SIZE = 64 * 1024
+
+        # Initialize gradient tensors
+        grad_intermediate_cache2 = torch.zeros_like(intermediate_cache2)
+        # Use float32 for grad_w2 accumulation to avoid bfloat16 precision loss
+        grad_w2 = torch.zeros(w2.shape, dtype=torch.float32, device=w2.device)
+        # Compute grad_topk_weights in DownProjFunction backward
+        grad_topk_weights = torch.zeros_like(topk_weights)
+
+        # Process in chunks to match forward pass
+        for chunk in range((num_tokens // CHUNK_SIZE) + 1):
+            begin_chunk_idx, end_chunk_idx = (
+                chunk * CHUNK_SIZE,
+                min((chunk + 1) * CHUNK_SIZE, num_tokens),
+            )
+
+            curr_num_tokens = end_chunk_idx - begin_chunk_idx
+            if curr_num_tokens == 0:
+                continue
+
+            curr_intermediate_cache2 = intermediate_cache2[begin_chunk_idx * topk : end_chunk_idx * topk]
+            curr_topk_ids = topk_ids[begin_chunk_idx:end_chunk_idx]
+            curr_grad_output = grad_output[begin_chunk_idx:end_chunk_idx]
+            curr_topk_weights = topk_weights[begin_chunk_idx:end_chunk_idx]
+
+            # 1. Calculate grad_intermediate_cache2: topk_weights * (grad_output @ w2)
+            for t in range(curr_num_tokens):
+                for k in range(topk):
+                    expert_id = curr_topk_ids[t, k].item()
+                    grad_y_tk = curr_grad_output[t, k]  # shape: (hidden_size,)
+                    W2_e = w2[expert_id]  # shape: (hidden_size, intermediate_size)
+                    weight_tk = curr_topk_weights[t, k]  # scalar
+
+                    grad_intermediate_cache2[(begin_chunk_idx + t) * topk + k] = weight_tk * (grad_y_tk @ W2_e)
+
+            # 2. Calculate grad_w2: topk_weights * (grad_output.T @ intermediate_cache2)
+            for t in range(curr_num_tokens):
+                for k in range(topk):
+                    expert_id = curr_topk_ids[t, k].item()
+                    grad_y_tk = curr_grad_output[t, k]  # shape: (hidden_size,)
+                    x_tk = curr_intermediate_cache2[t * topk + k]  # shape: (intermediate_size,)
+                    weight_tk = curr_topk_weights[t, k]  # scalar
+
+                    # Accumulate in float32
+                    grad_w2[expert_id] += (weight_tk * torch.outer(grad_y_tk, x_tk)).to(torch.float32)
+
+            # 3. Calculate grad_topk_weights: dot(grad_output, forward_output_before_weighting)
+            for t in range(curr_num_tokens):
+                for k in range(topk):
+                    expert_id = curr_topk_ids[t, k].item()
+                    grad_y_tk = curr_grad_output[t, k]  # shape: (hidden_size,)
+                    x_tk = curr_intermediate_cache2[t * topk + k]  # shape: (intermediate_size,)
+                    W2_e = w2[expert_id]  # shape: (hidden_size, intermediate_size)
+
+                    # Compute forward output before weighting
+                    forward_output_unweighted = x_tk @ W2_e.T  # shape: (hidden_size,)
+
+                    # grad_topk_weights: dot product
+                    grad_topk_weights[begin_chunk_idx + t, k] += torch.sum(grad_y_tk * forward_output_unweighted)
+
+        # Convert grad_w2 back to original dtype (bfloat16)
+        grad_w2 = grad_w2.to(intermediate_cache2.dtype)
+
+        return grad_intermediate_cache2, grad_w2, grad_topk_weights, None
+
+
+# ============================================================================
+# Import Triton Implementation
+# ============================================================================
+
+from slime.backends.fsdp_utils.kernels.fused_experts import DownProjFunction as DownProjFunctionTriton
+from slime.backends.fsdp_utils.kernels.fused_experts import GateUpProjFunction as GateUpProjFunctionTriton
+
+# ============================================================================
+# Test Fixtures and Utilities
+# ============================================================================
+
+
+@pytest.fixture
+def setup_moe_params():
+    """Setup MOE parameters for testing."""
+    torch.manual_seed(42)
+
+    # Small parameters for easier debugging
+    num_tokens = 64
+    hidden_size = 128
+    intermediate_size = 256
+    num_experts = 4
+    topk = 2
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.bfloat16
+
+    # Create input tensors with random values for better testing
+    hidden_states = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
+
+    # Create expert weights
+    w1 = torch.randn(num_experts, intermediate_size * 2, hidden_size, device=device, dtype=dtype)
+    w2 = torch.randn(num_experts, hidden_size, intermediate_size, device=device, dtype=dtype)
+
+    # Create router outputs
+    topk_weights = torch.rand(num_tokens, topk, device=device, dtype=dtype)
+    topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)  # normalize
+
+    # Random expert selection
+    topk_ids = torch.stack([torch.randperm(num_experts, device=device)[:topk] for _ in range(num_tokens)], dim=0).to(
+        torch.int32
+    )
+
+    return {
+        "hidden_states": hidden_states,
+        "w1": w1,
+        "w2": w2,
+        "topk_weights": topk_weights,
+        "topk_ids": topk_ids,
+        "device": device,
+        "dtype": dtype,
+    }
+
+
+# ============================================================================
+# Test Cases
+# ============================================================================
+
+
+class TestGateUpProjBackward:
+    """Test GateUpProjFunction backward pass comparison."""
+
+    def test_forward_consistency(self, setup_moe_params):
+        """Test that Triton and Python implementations produce same forward output."""
+        params = setup_moe_params
+
+        # Python implementation
+        out_python = GateUpProjFunctionPython.apply(
+            params["hidden_states"].clone(),
+            params["w1"].clone(),
+            params["topk_weights"].clone(),
+            params["topk_ids"].clone(),
+        )
+
+        # Triton implementation
+        out_triton = GateUpProjFunctionTriton.apply(
+            params["hidden_states"].clone(),
+            params["w1"].clone(),
+            params["topk_weights"].clone(),
+            params["topk_ids"].clone(),
+        )
+
+        # Check outputs are close
+        torch.testing.assert_close(out_python, out_triton, rtol=1, atol=1)
+        print("✓ GateUpProjFunction forward test passed")
+
+    def test_backward_consistency(self, setup_moe_params):
+        """Test that Triton and Python implementations produce same gradients."""
+        params = setup_moe_params
+
+        # Prepare inputs with requires_grad
+        hidden_states_python = params["hidden_states"].clone().requires_grad_(True)
+        w1_python = params["w1"].clone().requires_grad_(True)
+        topk_weights_python = params["topk_weights"].clone().requires_grad_(True)
+        topk_ids_python = params["topk_ids"].clone()
+
+        hidden_states_triton = params["hidden_states"].clone().requires_grad_(True)
+        w1_triton = params["w1"].clone().requires_grad_(True)
+        topk_weights_triton = params["topk_weights"].clone().requires_grad_(True)
+        topk_ids_triton = params["topk_ids"].clone()
+
+        # Python implementation
+        out_python = GateUpProjFunctionPython.apply(
+            hidden_states_python,
+            w1_python,
+            topk_weights_python,
+            topk_ids_python,
+        )
+
+        # Triton implementation
+        out_triton = GateUpProjFunctionTriton.apply(
+            hidden_states_triton,
+            w1_triton,
+            topk_weights_triton,
+            topk_ids_triton,
+        )
+
+        # Create gradient for backward
+        grad_output = torch.randn_like(out_python)
+
+        # Backward pass
+        out_python.backward(grad_output.clone())
+        out_triton.backward(grad_output.clone())
+
+        # Check hidden_states gradients
+        print("\n" + "=" * 80)
+        print("GateUpProjFunction Backward - hidden_states gradients:")
+        print("=" * 80)
+        if hidden_states_python.grad is not None and hidden_states_triton.grad is not None:
+            diff = hidden_states_python.grad - hidden_states_triton.grad
+            max_diff = torch.max(torch.abs(diff))
+            print(f"Max absolute difference: {max_diff:.6f}")
+            torch.testing.assert_close(hidden_states_python.grad, hidden_states_triton.grad, rtol=1, atol=1)
+            print("✓ hidden_states gradient matches")
+        print("=" * 80 + "\n")
+
+        # Check w1 gradients
+        print("\n" + "=" * 80)
+        print("GateUpProjFunction Backward - w1 gradients:")
+        print("=" * 80)
+        if w1_python.grad is not None and w1_triton.grad is not None:
+            diff = w1_python.grad - w1_triton.grad
+            max_diff = torch.max(torch.abs(diff))
+            print(f"Max absolute difference: {max_diff:.6f}")
+            torch.testing.assert_close(w1_python.grad, w1_triton.grad, rtol=1, atol=1)
+            print("✓ w1 gradient matches")
+        print("=" * 80 + "\n")
+
+        print("✓ GateUpProjFunction backward test passed")
+
+
+class TestDownProjBackward:
+    """Test DownProjFunction backward pass comparison."""
+
+    def test_forward_consistency(self, setup_moe_params):
+        """Test that Triton and Python implementations produce same forward output."""
+        params = setup_moe_params
+
+        # Create intermediate input (after SiluAndMul)
+        num_tokens = params["hidden_states"].shape[0]
+        topk = params["topk_ids"].shape[1]
+        intermediate_size = params["w2"].shape[2]
+        intermediate_cache2 = torch.randn(
+            num_tokens * topk, intermediate_size, device=params["device"], dtype=params["dtype"]
+        )
+
+        # Python implementation
+        out_python = DownProjFunctionPython.apply(
+            intermediate_cache2.clone(),
+            params["w2"].clone(),
+            params["topk_weights"].clone(),
+            params["topk_ids"].clone(),
+        )
+
+        # Triton implementation
+        out_triton = DownProjFunctionTriton.apply(
+            intermediate_cache2.clone(),
+            params["w2"].clone(),
+            params["topk_weights"].clone(),
+            params["topk_ids"].clone(),
+        )
+
+        # Check outputs are close
+        torch.testing.assert_close(out_python, out_triton, rtol=1, atol=1)
+        print("✓ DownProjFunction forward test passed")
+
+    def test_backward_consistency(self, setup_moe_params):
+        """Test that Triton and Python implementations produce same gradients."""
+        params = setup_moe_params
+
+        # Create intermediate input
+        num_tokens = params["hidden_states"].shape[0]
+        topk = params["topk_ids"].shape[1]
+        intermediate_size = params["w2"].shape[2]
+
+        intermediate_cache2_base = torch.randn(
+            num_tokens * topk, intermediate_size, device=params["device"], dtype=params["dtype"]
+        )
+
+        intermediate_cache2_python = intermediate_cache2_base.clone().requires_grad_(True)
+        intermediate_cache2_triton = intermediate_cache2_base.clone().requires_grad_(True)
+
+        w2_python = params["w2"].clone().requires_grad_(True)
+        w2_triton = params["w2"].clone().requires_grad_(True)
+
+        topk_weights_python = params["topk_weights"].clone().requires_grad_(True)
+        topk_weights_triton = params["topk_weights"].clone().requires_grad_(True)
+
+        # Python implementation
+        out_python = DownProjFunctionPython.apply(
+            intermediate_cache2_python,
+            w2_python,
+            topk_weights_python,
+            params["topk_ids"],
+        )
+
+        # Triton implementation
+        out_triton = DownProjFunctionTriton.apply(
+            intermediate_cache2_triton,
+            w2_triton,
+            topk_weights_triton,
+            params["topk_ids"],
+        )
+
+        # Create gradient for backward
+        grad_output = torch.randn_like(out_python)
+
+        # Backward pass
+        out_python.backward(grad_output.clone())
+        out_triton.backward(grad_output.clone())
+
+        # Check intermediate_cache2 gradients
+        print("\n" + "=" * 80)
+        print("DownProjFunction Backward - intermediate_cache2 gradients:")
+        print("=" * 80)
+        if intermediate_cache2_python.grad is not None and intermediate_cache2_triton.grad is not None:
+            diff = intermediate_cache2_python.grad - intermediate_cache2_triton.grad
+            max_diff = torch.max(torch.abs(diff))
+            print(f"Max absolute difference: {max_diff:.6f}")
+            torch.testing.assert_close(
+                intermediate_cache2_python.grad, intermediate_cache2_triton.grad, rtol=1, atol=1
+            )
+            print("✓ intermediate_cache2 gradient matches")
+        print("=" * 80 + "\n")
+
+        # Check topk_weights gradients
+        print("\n" + "=" * 80)
+        print("DownProjFunction Backward - topk_weights gradients:")
+        print("=" * 80)
+        if topk_weights_python.grad is not None and topk_weights_triton.grad is not None:
+            diff = topk_weights_python.grad - topk_weights_triton.grad
+            max_diff = torch.max(torch.abs(diff))
+            print(f"Max absolute difference: {max_diff:.6f}")
+            torch.testing.assert_close(topk_weights_python.grad, topk_weights_triton.grad, rtol=1, atol=1)
+            print("✓ topk_weights gradient matches")
+        print("=" * 80 + "\n")
+
+        # Check w2 gradients
+        print("\n" + "=" * 80)
+        print("DownProjFunction Backward - w2 gradients:")
+        print("=" * 80)
+        if w2_python.grad is not None and w2_triton.grad is not None:
+            diff = w2_python.grad - w2_triton.grad
+            max_diff = torch.max(torch.abs(diff))
+            print(f"Max absolute difference: {max_diff:.6f}")
+            torch.testing.assert_close(w2_python.grad, w2_triton.grad, rtol=1, atol=1)
+            print("✓ w2 gradient matches")
+        print("=" * 80 + "\n")
+
+        print("✓ DownProjFunction backward test passed")
+
+
+# ============================================================================
+# Main Test Runner
+# ============================================================================
+
+
+def run_all_tests():
+    """Run all tests."""
+    print("=" * 80)
+    print("Running Fused Experts Backward Tests")
+    print("Testing: Triton Implementation vs Python Reference")
+    print("=" * 80)
+
+    if not torch.cuda.is_available():
+        print("WARNING: CUDA not available, skipping tests")
+        return
+
+    # Setup parameters
+    torch.manual_seed(42)
+    params_dict = {}
+
+    # Small parameters for testing
+    num_tokens = 64
+    hidden_size = 128
+    intermediate_size = 256
+    num_experts = 4
+    topk = 2
+
+    device = "cuda"
+    dtype = torch.bfloat16
+
+    # Create input tensors
+    params_dict["hidden_states"] = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
+    params_dict["w1"] = torch.randn(num_experts, intermediate_size * 2, hidden_size, device=device, dtype=dtype)
+    params_dict["w2"] = torch.randn(num_experts, hidden_size, intermediate_size, device=device, dtype=dtype)
+    params_dict["topk_weights"] = torch.rand(num_tokens, topk, device=device, dtype=dtype)
+    params_dict["topk_weights"] = params_dict["topk_weights"] / params_dict["topk_weights"].sum(dim=-1, keepdim=True)
+    params_dict["topk_ids"] = torch.stack(
+        [torch.randperm(num_experts, device=device)[:topk] for _ in range(num_tokens)], dim=0
+    ).to(torch.int32)
+    params_dict["device"] = device
+    params_dict["dtype"] = dtype
+
+    print("\n" + "=" * 80)
+    print("Testing GateUpProjFunction Backward")
+    print("=" * 80)
+    test_gate_up = TestGateUpProjBackward()
+    test_gate_up.test_forward_consistency(params_dict)
+    test_gate_up.test_backward_consistency(params_dict)
+
+    print("\n" + "=" * 80)
+    print("Testing DownProjFunction Backward")
+    print("=" * 80)
+    test_down = TestDownProjBackward()
+    test_down.test_forward_consistency(params_dict)
+    test_down.test_backward_consistency(params_dict)
+
+    print("\n" + "=" * 80)
+    print("All Backward Tests Passed! ✓")
+    print("=" * 80)
+
+
+if __name__ == "__main__":
+    run_all_tests()
diff --git a/tests/test_glm4.7_30B_A3B_pd_mooncake.py b/tests/test_glm4.7_30B_A3B_pd_mooncake.py
deleted file mode 100644
index 8ea17b5ce9..0000000000
--- a/tests/test_glm4.7_30B_A3B_pd_mooncake.py
+++ /dev/null
@@ -1,160 +0,0 @@
-"""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-speculative-algorithm EAGLE "
-        "--sglang-speculative-num-steps 3 "
-        "--sglang-speculative-eagle-topk 1 "
-        "--sglang-speculative-num-draft-tokens 4 "
-        "--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_gspo.sh b/tests/test_gspo.sh
index d26a479038..6e915ca652 100644
--- a/tests/test_gspo.sh
+++ b/tests/test_gspo.sh
@@ -71,7 +71,7 @@ ray job submit --address="http://127.0.0.1:8265" \
    --actor-num-nodes 1 \
    --actor-num-gpus-per-node 4 \
    --colocate \
-   --train-backend megatron \
+   --train-backend fsdp \
    ${CKPT_ARGS[@]} \
    ${ROLLOUT_ARGS[@]} \
    ${OPTIMIZER_ARGS[@]} \
diff --git a/tests/test_megatron_argument_validation.py b/tests/test_megatron_argument_validation.py
deleted file mode 100644
index 83e8390ee1..0000000000
--- a/tests/test_megatron_argument_validation.py
+++ /dev/null
@@ -1,140 +0,0 @@
-import importlib.util
-import sys
-import types
-from pathlib import Path
-
-import pytest
-
-
-def load_arguments_module(monkeypatch):
-    megatron_mod = types.ModuleType("megatron")
-    training_mod = types.ModuleType("megatron.training")
-    arguments_mod = types.ModuleType("megatron.training.arguments")
-    tokenizer_pkg_mod = types.ModuleType("megatron.training.tokenizer")
-    tokenizer_mod = types.ModuleType("megatron.training.tokenizer.tokenizer")
-    transformers_mod = types.ModuleType("transformers")
-
-    arguments_mod.parse_args = lambda *args, **kwargs: None
-    arguments_mod.validate_args = lambda args: args
-    tokenizer_mod._vocab_size_with_padding = lambda vocab_size, _args: vocab_size
-    transformers_mod.AutoConfig = types.SimpleNamespace(from_pretrained=lambda *args, **kwargs: None)
-
-    monkeypatch.setitem(sys.modules, "megatron", megatron_mod)
-    monkeypatch.setitem(sys.modules, "megatron.training", training_mod)
-    monkeypatch.setitem(sys.modules, "megatron.training.arguments", arguments_mod)
-    monkeypatch.setitem(sys.modules, "megatron.training.tokenizer", tokenizer_pkg_mod)
-    monkeypatch.setitem(sys.modules, "megatron.training.tokenizer.tokenizer", tokenizer_mod)
-    monkeypatch.setitem(sys.modules, "transformers", transformers_mod)
-
-    module_path = Path(__file__).resolve().parents[1] / "slime" / "backends" / "megatron_utils" / "arguments.py"
-    module_name = "test_megatron_argument_validation_module"
-    sys.modules.pop(module_name, None)
-    spec = importlib.util.spec_from_file_location(module_name, module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
-def make_qwen3_6_args(**overrides):
-    values = dict(
-        hidden_size=2048,
-        num_attention_heads=16,
-        num_layers=40,
-        ffn_hidden_size=512,
-        moe_ffn_hidden_size=512,
-        moe_shared_expert_intermediate_size=512,
-        moe_layer_freq=[1] * 40,
-        untie_embeddings_and_output_weights=True,
-        norm_epsilon=1e-6,
-        layernorm_epsilon=1e-6,
-        rotary_base=10000000,
-    )
-    values.update(overrides)
-    return types.SimpleNamespace(**values)
-
-
-def make_qwen3_6_hf_config():
-    text_config = types.SimpleNamespace(
-        hidden_size=2048,
-        num_attention_heads=16,
-        num_hidden_layers=40,
-        intermediate_size=5632,
-        moe_intermediate_size=512,
-        shared_expert_intermediate_size=512,
-        num_experts=256,
-        tie_word_embeddings=False,
-        rms_norm_eps=1e-6,
-        rope_parameters={"rope_theta": 10000000},
-    )
-    return types.SimpleNamespace(text_config=text_config)
-
-
-def make_allgather_cp_args(**overrides):
-    values = dict(
-        allgather_cp=True,
-        context_parallel_size=2,
-    )
-    values.update(overrides)
-    return types.SimpleNamespace(**values)
-
-
-@pytest.mark.unit
-def test_hf_validate_all_moe_skips_dense_intermediate_size(monkeypatch):
-    module = load_arguments_module(monkeypatch)
-
-    module._hf_validate_args(make_qwen3_6_args(), make_qwen3_6_hf_config())
-
-
-@pytest.mark.unit
-def test_hf_validate_checks_moe_intermediate_size(monkeypatch):
-    module = load_arguments_module(monkeypatch)
-
-    with pytest.raises(AssertionError, match="moe_intermediate_size"):
-        module._hf_validate_args(make_qwen3_6_args(moe_ffn_hidden_size=256), make_qwen3_6_hf_config())
-
-
-@pytest.mark.unit
-def test_hf_validate_checks_dense_intermediate_size_when_moe_has_dense_layers(monkeypatch):
-    module = load_arguments_module(monkeypatch)
-
-    args = make_qwen3_6_args(moe_layer_freq=[0] + [1] * 39)
-
-    with pytest.raises(AssertionError, match="intermediate_size"):
-        module._hf_validate_args(args, make_qwen3_6_hf_config())
-
-
-@pytest.mark.unit
-def test_allgather_cp_rejects_non_dsa_cp_models(monkeypatch):
-    module = load_arguments_module(monkeypatch)
-    args = make_allgather_cp_args()
-    hf_config = types.SimpleNamespace(architectures=["Qwen3ForCausalLM"], model_type="qwen3")
-
-    with pytest.raises(ValueError, match="only supported for DSA attention models"):
-        module._validate_allgather_cp_supported(args, hf_config)
-
-
-@pytest.mark.unit
-@pytest.mark.parametrize(
-    "hf_config",
-    [
-        types.SimpleNamespace(architectures=["DeepseekV32ForCausalLM"], model_type="deepseek_v3"),
-        types.SimpleNamespace(architectures=["GlmMoeDsaForCausalLM"], model_type="glm"),
-    ],
-)
-def test_allgather_cp_allows_dsa_architectures(monkeypatch, hf_config):
-    module = load_arguments_module(monkeypatch)
-
-    module._validate_allgather_cp_supported(make_allgather_cp_args(), hf_config)
-
-
-@pytest.mark.unit
-def test_allgather_cp_ignores_cp_size_one(monkeypatch):
-    module = load_arguments_module(monkeypatch)
-    args = make_allgather_cp_args(context_parallel_size=1)
-
-    module._validate_allgather_cp_supported(args)
-
-
-if __name__ == "__main__":
-    raise SystemExit(pytest.main([__file__]))
diff --git a/tests/test_moonlight_16B_A3B_r3.py b/tests/test_moonlight_16B_A3B_r3.py
index facb2d0b70..b8a4e5ee3a 100644
--- a/tests/test_moonlight_16B_A3B_r3.py
+++ b/tests/test_moonlight_16B_A3B_r3.py
@@ -69,6 +69,7 @@ def execute():
         "--entropy-coef 0.00 "
         "--eps-clip 4e-4 "
         "--use-rollout-routing-replay "
+        "--use-slime-router "
     )
 
     optimizer_args = (
diff --git a/tests/test_quick_start_glm4_9B.py b/tests/test_quick_start_glm4_9B.py
index 30cb348594..4f87df8eab 100644
--- a/tests/test_quick_start_glm4_9B.py
+++ b/tests/test_quick_start_glm4_9B.py
@@ -80,7 +80,7 @@ def execute():
         "--adam-beta2 0.98 "
     )
 
-    sglang_args = "--rollout-num-gpus-per-engine 2 " "--sglang-cuda-graph-max-bs 32 "
+    sglang_args = "--rollout-num-gpus-per-engine 2 " "--sglang-cuda-graph-max-bs 32 " "--use-slime-router "
 
     ci_args = "--ci-test "
 
diff --git a/tests/test_qwen2.5_0.5B_async_short.py b/tests/test_qwen2.5_0.5B_async_short.py
deleted file mode 100644
index c3925d4432..0000000000
--- a/tests/test_qwen2.5_0.5B_async_short.py
+++ /dev/null
@@ -1,121 +0,0 @@
-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"
-NUM_GPUS = 4
-
-
-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")
-
-
-def execute():
-    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
-
-    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 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 8192 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-        "--balance-data "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 1 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 9216 "
-    )
-
-    grpo_args = (
-        "--advantage-estimator grpo "
-        "--use-kl-loss "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type low_var_kl "
-        "--entropy-coef 0.00 "
-        "--eps-clip 0.2 "
-        "--eps-clip-high 0.28 "
-    )
-
-    optimizer_args = (
-        "--optimizer adam "
-        "--lr 1e-6 "
-        "--lr-decay-style constant "
-        "--weight-decay 0.1 "
-        "--adam-beta1 0.9 "
-        "--adam-beta2 0.98 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 1 "
-        f"--sglang-mem-fraction-static {0.55 if TIGHT_DEVICE_MEMORY else 0.65} "
-        "--sglang-cuda-graph-max-bs 32 "
-        "--sglang-enable-metrics "
-    )
-
-    ci_args = "--ci-test "
-
-    fault_tolerance_args = (
-        "--use-fault-tolerance "
-        "--rollout-health-check-interval 5 "
-        "--rollout-health-check-timeout 10 "
-        "--rollout-health-check-first-wait 0 "
-    )
-
-    misc_args = (
-        "--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 1 "
-        "--rollout-num-gpus 3 "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    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"{ci_args} "
-        f"{fault_tolerance_args} "
-        f"{misc_args} "
-    )
-
-    U.execute_train(
-        train_args=train_args,
-        num_gpus_per_node=NUM_GPUS,
-        megatron_model_type=MODEL_TYPE,
-        train_script="train_async.py",
-    )
-
-
-if __name__ == "__main__":
-    prepare()
-    os.environ.pop("http_proxy")
-    os.environ.pop("https_proxy")
-    os.environ.pop("HTTP_PROXY")
-    os.environ.pop("HTTPS_PROXY")
-    execute()
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 6e687799c8..63fd24904a 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"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.exec_command(f"huggingface-cli download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
     U.hf_download_dataset("zhuzilin/gsm8k")
 
 
diff --git a/tests/test_qwen3.5_0.8B_gsm8k_async_short.py b/tests/test_qwen2.5_0.5B_gsm8k_async_short.py
similarity index 85%
rename from tests/test_qwen3.5_0.8B_gsm8k_async_short.py
rename to tests/test_qwen2.5_0.5B_gsm8k_async_short.py
index 14d92ff4fe..ee71ff60e6 100644
--- a/tests/test_qwen3.5_0.8B_gsm8k_async_short.py
+++ b/tests/test_qwen2.5_0.5B_gsm8k_async_short.py
@@ -1,30 +1,21 @@
 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 = "Qwen3.5-0.8B"
-MODEL_TYPE = "qwen3.5-0.8B"
+MODEL_NAME = "Qwen2.5-0.5B-Instruct"
+MODEL_TYPE = "qwen2.5-0.5B"
 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"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.exec_command(f"huggingface-cli 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 {TORCH_DIST_CKPT} "
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
 
     rollout_args = (
         "--prompt-data /root/datasets/gsm8k/train.parquet "
@@ -103,10 +94,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 = (
@@ -133,6 +124,8 @@ def execute():
 
 if __name__ == "__main__":
     prepare()
-    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
-        os.environ.pop(proxy_var, None)
+    os.environ.pop("http_proxy")
+    os.environ.pop("https_proxy")
+    os.environ.pop("HTTP_PROXY")
+    os.environ.pop("HTTPS_PROXY")
     execute()
diff --git a/tests/test_qwen3.5_0.8B_gsm8k_short.py b/tests/test_qwen2.5_0.5B_gsm8k_short.py
similarity index 85%
rename from tests/test_qwen3.5_0.8B_gsm8k_short.py
rename to tests/test_qwen2.5_0.5B_gsm8k_short.py
index 856413de75..1b47c8d007 100644
--- a/tests/test_qwen3.5_0.8B_gsm8k_short.py
+++ b/tests/test_qwen2.5_0.5B_gsm8k_short.py
@@ -1,30 +1,21 @@
 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 = "Qwen3.5-0.8B"
-MODEL_TYPE = "qwen3.5-0.8B"
+MODEL_NAME = "Qwen2.5-0.5B-Instruct"
+MODEL_TYPE = "qwen2.5-0.5B"
 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"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.exec_command(f"huggingface-cli 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 {TORCH_DIST_CKPT} "
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
 
     rollout_args = (
         "--prompt-data /root/datasets/gsm8k/train.parquet "
@@ -103,10 +94,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 = (
@@ -132,6 +123,8 @@ def execute():
 
 if __name__ == "__main__":
     prepare()
-    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
-        os.environ.pop(proxy_var, None)
+    os.environ.pop("http_proxy")
+    os.environ.pop("https_proxy")
+    os.environ.pop("HTTP_PROXY")
+    os.environ.pop("HTTPS_PROXY")
     execute()
diff --git a/tests/test_qwen2.5_0.5B_opd_sglang.py b/tests/test_qwen2.5_0.5B_opd_sglang.py
index eb0892fb91..85addd0ab4 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"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}")
+    U.exec_command(f"huggingface-cli 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
deleted file mode 100644
index c03d863d8e..0000000000
--- a/tests/test_qwen2.5_0.5B_ppo_critic_only_short.py
+++ /dev/null
@@ -1,131 +0,0 @@
-import os
-import tempfile
-
-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"
-NUM_GPUS = 4
-
-
-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")
-
-
-def execute():
-    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
-    megatron_config.write(
-        """
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-"""
-    )
-    megatron_config.close()
-
-    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ "
-
-    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 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 1024 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-        "--balance-data "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 1 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 9216 "
-    )
-
-    ppo_args = (
-        "--advantage-estimator ppo "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type k1 "
-        "--kl-coef 0.00 "
-        "--entropy-coef 0.00 "
-        "--eps-clip 4e-4 "
-        "--num-critic-only-steps 3 "
-        "--normalize-advantages "
-    )
-
-    optimizer_args = (
-        "--optimizer adam "
-        "--lr 1e-6 "
-        "--lr-decay-style constant "
-        "--weight-decay 0.1 "
-        "--adam-beta1 0.9 "
-        "--adam-beta2 0.98 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 1 "
-        "--rollout-num-gpus 2 "
-        f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
-        "--sglang-cuda-graph-max-bs 32 "
-        "--sglang-enable-metrics "
-    )
-
-    ci_args = "--ci-test "
-
-    misc_args = (
-        "--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 4 "
-        "--megatron-to-hf-mode bridge "
-        "--colocate "
-    )
-
-    train_args = (
-        f"--megatron-config-path {megatron_config.name} "
-        f"{ckpt_args} "
-        f"{rollout_args} "
-        f"{optimizer_args} "
-        f"{ppo_args} "
-        f"{U.get_default_wandb_args(__file__)} "
-        f"{perf_args} "
-        f"{sglang_args} "
-        f"{ci_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_sglang_config.py b/tests/test_qwen2.5_0.5B_sglang_config.py
index f30c6ac6cf..0b1ec9e057 100644
--- a/tests/test_qwen2.5_0.5B_sglang_config.py
+++ b/tests/test_qwen2.5_0.5B_sglang_config.py
@@ -16,7 +16,7 @@
 SGLANG_CONFIG_YAML = """\
 sglang:
   - name: default
-    server_groups:
+    engine_groups:
       - worker_type: regular
         num_gpus: 4
         num_gpus_per_engine: 2
@@ -30,7 +30,7 @@
 
 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.exec_command(f"huggingface-cli 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 68215b34c6..41d70bea44 100644
--- a/tests/test_qwen2.5_0.5B_sglang_config_distributed.py
+++ b/tests/test_qwen2.5_0.5B_sglang_config_distributed.py
@@ -10,14 +10,14 @@
 NUM_GPUS = 8
 
 # Inline sglang config: same model, 3 engine groups with different parallelism.
-# Non-colocated (decoupled training and rollout): actor uses 4 GPUs, rollout uses 4 GPUs.
+# Non-colocated (训推分离): actor uses 4 GPUs, rollout uses 4 GPUs.
 # Group 1: 2 GPUs, 2 GPUs/engine (tp=2) → 1 engine
 # Group 2: 1 GPU,  1 GPU/engine  (tp=1) → 1 engine
 # Group 3: 1 GPU,  placeholder   → reserves 1 GPU slot, no engine created
 SGLANG_CONFIG_YAML = """\
 sglang:
   - name: default
-    server_groups:
+    engine_groups:
       - worker_type: regular
         num_gpus: 2
         num_gpus_per_engine: 2
@@ -31,7 +31,7 @@
 
 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.exec_command(f"huggingface-cli 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
deleted file mode 100644
index 6f45095bfb..0000000000
--- a/tests/test_qwen2.5_0.5B_short.py
+++ /dev/null
@@ -1,120 +0,0 @@
-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"
-NUM_GPUS = 4
-
-
-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")
-
-
-def execute():
-    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
-
-    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 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 8192 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-        "--balance-data "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 1 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 9216 "
-    )
-
-    grpo_args = (
-        "--advantage-estimator grpo "
-        "--use-kl-loss "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type low_var_kl "
-        "--entropy-coef 0.00 "
-        "--eps-clip 0.2 "
-        "--eps-clip-high 0.28 "
-    )
-
-    optimizer_args = (
-        "--optimizer adam "
-        "--lr 1e-6 "
-        "--lr-decay-style constant "
-        "--weight-decay 0.1 "
-        "--adam-beta1 0.9 "
-        "--adam-beta2 0.98 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 1 "
-        f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
-        "--sglang-cuda-graph-max-bs 32 "
-        "--sglang-enable-metrics "
-    )
-
-    ci_args = "--ci-test "
-
-    fault_tolerance_args = (
-        "--use-fault-tolerance "
-        "--rollout-health-check-interval 5 "
-        "--rollout-health-check-timeout 10 "
-        "--rollout-health-check-first-wait 0 "
-    )
-
-    misc_args = (
-        "--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 4 "
-        "--colocate "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    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"{ci_args} "
-        f"{fault_tolerance_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()
-    os.environ.pop("http_proxy")
-    os.environ.pop("https_proxy")
-    os.environ.pop("HTTP_PROXY")
-    os.environ.pop("HTTPS_PROXY")
-    execute()
diff --git a/tests/test_qwen3.6_35B_A3B_pd_mooncake.py b/tests/test_qwen3.6_35B_A3B_pd_mooncake.py
deleted file mode 100644
index 83e2449dae..0000000000
--- a/tests/test_qwen3.6_35B_A3B_pd_mooncake.py
+++ /dev/null
@@ -1,153 +0,0 @@
-import os
-import tempfile
-
-import slime.utils.external_utils.command_utils as U
-
-
-MODEL_NAME = "Qwen3.6-35B-A3B"
-MODEL_TYPE = "qwen3.5-35B-A3B"
-NUM_GPUS = 8
-TORCH_DIST_CKPT = f"/root/models/{MODEL_NAME}_torch_dist"
-
-
-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,
-        dir_dst="/root/models",
-    )
-
-
-def execute():
-    debug_data_path = os.environ.get("DEBUG_ROLLOUT_DATA") or tempfile.mktemp(
-        prefix="qwen3_6_35b_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 {TORCH_DIST_CKPT} "
-
-    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 2 "
-        "--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-disaggregation-transfer-backend mooncake "
-        "--sglang-speculative-algorithm EAGLE "
-        "--sglang-speculative-num-steps 3 "
-        "--sglang-speculative-eagle-topk 1 "
-        "--sglang-speculative-num-draft-tokens 4 "
-        "--sglang-mamba-scheduler-strategy extra_buffer "
-        "--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_0.6B_megatron_fsdp_align.py b/tests/test_qwen3_0.6B_megatron_fsdp_align.py
new file mode 100644
index 0000000000..1c48f42a69
--- /dev/null
+++ b/tests/test_qwen3_0.6B_megatron_fsdp_align.py
@@ -0,0 +1,154 @@
+import os
+
+import slime.utils.external_utils.command_utils as U
+
+MODEL_NAME = "Qwen3-0.6B"
+MODEL_TYPE = "qwen3-0.6B"
+NUM_GPUS = 4
+CP_SIZE = 1
+MEGATRON_TP_SIZE = 1
+MEGATRON_PP_SIZE = 1
+
+
+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.convert_checkpoint(
+        model_name=MODEL_NAME,
+        megatron_model_type=MODEL_TYPE,
+        num_gpus_per_node=NUM_GPUS,
+        dir_dst="/root/models",
+    )
+
+
+def execute():
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/"
+
+    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 1 "
+        "--rollout-batch-size 8 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 8192 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 64 "
+        "--use-dynamic-batch-size "
+        "--max-tokens-per-gpu 8192 "
+    )
+
+    ppo_args = (
+        "--advantage-estimator grpo "
+        "--kl-loss-coef 0.00 "
+        "--kl-loss-type k1 "
+        "--kl-coef 0.00 "
+        "--entropy-coef 0.00 "
+        "--eps-clip 4e-4 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--sglang-chunked-prefill-size 4096 "
+        "--sglang-mem-fraction-static 0.75 "
+        "--sglang-cuda-graph-max-bs 32 "
+    )
+
+    ci_args = "--ci-test "
+
+    misc_args = "--actor-num-nodes 1 " "--colocate " f"--actor-num-gpus-per-node {NUM_GPUS} "
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{ppo_args} "
+        f"{U.get_default_wandb_args(__file__)} "
+        f"{sglang_args} "
+        f"{ci_args} "
+        f"{misc_args} "
+    )
+
+    debug_data_path = "test_rollout_data_megatron_fsdp_align.pt"
+    grad_norm_path = "grad_norm_fsdp.pt"
+
+    fsdp_args = (
+        "--train-backend fsdp "
+        "--attn-implementation flash_attention_2 "
+        "--gradient-checkpointing "
+        f"--update-weight-buffer-size {512 * 1024 * 1024} "
+        """--train-env-vars '{"PYTORCH_CUDA_ALLOC_CONF":"expandable_segments:True"}' """
+    )
+
+    try:
+        U.execute_train(
+            train_args=train_args + (f"{fsdp_args}" f"--save-debug-rollout-data {debug_data_path} "),
+            num_gpus_per_node=NUM_GPUS,
+            megatron_model_type=None,
+        )
+
+        U.execute_train(
+            train_args=train_args
+            + (
+                f"{fsdp_args}"
+                f"--load-debug-rollout-data {debug_data_path} "
+                f"--ci-save-grad-norm {grad_norm_path} "
+                "--debug-train-only "
+            ),
+            num_gpus_per_node=NUM_GPUS,
+            megatron_model_type=None,
+        )
+
+        U.execute_train(
+            train_args=train_args
+            + (
+                f"--ref-load /root/models/{MODEL_NAME}_torch_dist "
+                f"--tensor-model-parallel-size {MEGATRON_TP_SIZE} "
+                "--sequence-parallel "
+                f"--pipeline-model-parallel-size {MEGATRON_PP_SIZE} "
+                f"--context-parallel-size {CP_SIZE} "
+                "--expert-model-parallel-size 1 "
+                "--expert-tensor-parallel-size 1 "
+                "--recompute-granularity full "
+                "--recompute-method uniform "
+                "--recompute-num-layers 1 "
+                "--train-memory-margin-bytes 3221225472 "
+                f"--load-debug-rollout-data {debug_data_path} "
+                f"--ci-load-grad-norm {grad_norm_path} "
+                "--attention-dropout 0.0 "
+                "--hidden-dropout 0.0 "
+                "--accumulate-allreduce-grads-in-fp32 "
+                "--attention-softmax-in-fp32 "
+                "--attention-backend flash "
+                "--debug-train-only "
+            ),
+            num_gpus_per_node=NUM_GPUS,
+            megatron_model_type=MODEL_TYPE,
+        )
+
+    finally:
+        if os.path.exists(grad_norm_path):
+            os.remove(grad_norm_path)
+        if os.path.exists(debug_data_path):
+            os.remove(debug_data_path)
+
+
+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_30B_A3B_r3.py b/tests/test_qwen3_30B_A3B_r3.py
index 71d87a3919..9beeb57686 100644
--- a/tests/test_qwen3_30B_A3B_r3.py
+++ b/tests/test_qwen3_30B_A3B_r3.py
@@ -77,6 +77,7 @@ def execute():
         "--eps-clip 4e-4 "
         "--use-tis "
         "--use-rollout-routing-replay "
+        "--use-slime-router "
     )
 
     optimizer_args = (
diff --git a/tests/test_qwen3_4B_fsdp_true_on_policy.py b/tests/test_qwen3_4B_fsdp_true_on_policy.py
new file mode 100644
index 0000000000..b365ddf665
--- /dev/null
+++ b/tests/test_qwen3_4B_fsdp_true_on_policy.py
@@ -0,0 +1,123 @@
+import os
+from argparse import ArgumentParser
+import slime.utils.external_utils.command_utils as U
+
+ENABLE_EVAL = bool(int(os.environ.get("SLIME_TEST_ENABLE_EVAL", "1")))
+NUM_GPUS = 4
+
+MODEL_NAME = "Qwen3-4B"
+
+parser = ArgumentParser()
+parser.add_argument("--colocated", action="store_true", help="Whether to run with colocate.")
+
+
+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")
+
+
+def execute(args):
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} "
+
+    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 math "
+        "--num-rollout 3 "
+        "--rollout-batch-size 8 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 64 "
+    )
+
+    eval_args = (
+        f"{'--eval-interval 20 ' if ENABLE_EVAL else ''}"
+        "--eval-prompt-data aime /root/datasets/aime-2024/aime-2024.jsonl "
+        "--n-samples-per-eval-prompt 1 "
+        "--eval-max-response-len 4096 "
+        "--eval-top-p 0.7 "
+    )
+
+    fsdp_args = "--train-backend fsdp " "--update-weight-buffer-size 536870912 "
+
+    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 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--sglang-decode-log-interval 1000 "
+        "--sglang-cuda-graph-max-bs 32 "
+        "--sglang-enable-metrics "
+        "--sglang-enable-deterministic-inference "
+        "--sglang-rl-on-policy-target fsdp "
+        "--sglang-attention-backend fa3 "
+        "--attn-implementation flash_attention_3 "
+        "--deterministic-mode "
+        "--true-on-policy-mode "
+    )
+
+    ci_args = "--ci-test "
+
+    if args.colocated:
+        misc_args = f"--actor-num-nodes 1 --actor-num-gpus-per-node {NUM_GPUS} --colocate "
+    else:
+        misc_args = (
+            f"--actor-num-nodes 1 --actor-num-gpus-per-node {NUM_GPUS // 2} --rollout-num-gpus {NUM_GPUS // 2} "
+        )
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{U.get_default_wandb_args(__file__)} "
+        f"{fsdp_args} "
+        f"{eval_args} "
+        f"{sglang_args} "
+        f"{ci_args} "
+        f"{misc_args} "
+    )
+
+    extra_env_vars = {
+        "NCCL_ALGO": "allreduce:tree",
+        "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
+        "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+    }
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=NUM_GPUS,
+        megatron_model_type=None,
+        extra_env_vars=extra_env_vars,
+    )
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    prepare()
+    for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"):
+        os.environ.pop(proxy_var, None)
+    execute(args)
diff --git a/tests/test_qwen3_4B_ppo.py b/tests/test_qwen3_4B_ppo.py
index 9b9105f499..b800c19fed 100644
--- a/tests/test_qwen3_4B_ppo.py
+++ b/tests/test_qwen3_4B_ppo.py
@@ -1,5 +1,4 @@
 import os
-import tempfile
 
 import slime.utils.external_utils.command_utils as U
 
@@ -22,22 +21,6 @@ def prepare():
 
 
 def execute():
-    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
-    megatron_config.write(
-        """
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-"""
-    )
-    megatron_config.close()
-
     ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/{MODEL_NAME}_torch_dist "
 
     rollout_args = (
@@ -86,6 +69,7 @@ def execute():
         "--eps-clip 4e-4 "
         "--num-critic-only-steps 1 "
         "--normalize-advantages "
+        "--critic-lr 1e-5 "
     )
 
     optimizer_args = (
@@ -118,12 +102,11 @@ def execute():
         # need to comment this when using model with MLA
         "--attention-backend flash "
         "--actor-num-nodes 1 "
-        "--actor-num-gpus-per-node 8 "
+        "--actor-num-gpus-per-node 4 "
         "--colocate "
     )
 
     train_args = (
-        f"--megatron-config-path {megatron_config.name} "
         f"{ckpt_args} "
         f"{rollout_args} "
         f"{optimizer_args} "
diff --git a/tests/test_qwen3_4B_ppo_disaggregate.py b/tests/test_qwen3_4B_ppo_disaggregate.py
deleted file mode 100644
index 52119a5b67..0000000000
--- a/tests/test_qwen3_4B_ppo_disaggregate.py
+++ /dev/null
@@ -1,150 +0,0 @@
-import os
-import tempfile
-
-import slime.utils.external_utils.command_utils as U
-
-
-ENABLE_EVAL = bool(int(os.environ.get("SLIME_TEST_ENABLE_EVAL", "1")))
-TIGHT_HOST_MEMORY = bool(int(os.environ.get("SLIME_TEST_TIGHT_HOST_MEMORY", "1")))
-
-MODEL_NAME = "Qwen3-4B"
-MODEL_TYPE = "qwen3-4B"
-NUM_GPUS = 8
-
-
-def prepare():
-    U.exec_command("mkdir -p /root/models /root/datasets")
-    U.exec_command("hf download Qwen/Qwen3-4B --local-dir /root/models/Qwen3-4B")
-    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():
-    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
-    megatron_config.write(
-        """
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-"""
-    )
-    megatron_config.close()
-
-    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 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 8192 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-        "--balance-data "
-    )
-
-    eval_args = (
-        f"{'--eval-interval 20 ' if ENABLE_EVAL else ''}"
-        "--eval-prompt-data aime24 /root/datasets/aime-2024/aime-2024.jsonl "
-        "--n-samples-per-eval-prompt 1 "
-        "--eval-max-response-len 16384 "
-        "--eval-top-k 1 "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 2 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 2 "
-        "--recompute-granularity full "
-        "--recompute-method uniform "
-        "--recompute-num-layers 1 "
-        "--use-dynamic-batch-size "
-        f"--max-tokens-per-gpu {2048 if TIGHT_HOST_MEMORY else 16384} "
-    )
-
-    ppo_args = (
-        "--advantage-estimator ppo "
-        f"{'' if TIGHT_HOST_MEMORY else '--use-kl-loss '}"
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type k1 "
-        "--kl-coef 0.00 "
-        "--entropy-coef 0.00 "
-        "--eps-clip 4e-4 "
-        "--num-critic-only-steps 1 "
-        "--normalize-advantages "
-    )
-
-    optimizer_args = (
-        "--optimizer adam "
-        "--lr 1e-6 "
-        "--lr-decay-style constant "
-        "--weight-decay 0.1 "
-        "--adam-beta1 0.9 "
-        "--adam-beta2 0.98 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 2 "
-        "--rollout-num-gpus 4 "
-        "--sglang-mem-fraction-static 0.8 "
-        "--sglang-cuda-graph-max-bs 32 "
-        "--sglang-max-running-requests 512 "
-        "--sglang-enable-metrics "
-    )
-
-    ci_args = "--ci-test "
-
-    misc_args = (
-        # default dropout in megatron is 0.1
-        "--attention-dropout 0.0 "
-        "--hidden-dropout 0.0 "
-        # should be good for model performance
-        "--accumulate-allreduce-grads-in-fp32 "
-        "--attention-softmax-in-fp32 "
-        # need to comment this when using model with MLA
-        "--attention-backend flash "
-        "--actor-num-nodes 1 "
-        "--actor-num-gpus-per-node 4 "
-    )
-
-    train_args = (
-        f"--megatron-config-path {megatron_config.name} "
-        f"{ckpt_args} "
-        f"{rollout_args} "
-        f"{optimizer_args} "
-        f"{ppo_args} "
-        f"{U.get_default_wandb_args(__file__)} "
-        f"{perf_args} "
-        f"{eval_args} "
-        f"{sglang_args} "
-        f"{ci_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__":
-    # TODO also use typer
-    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_4B_ppo_train_critic_only.py b/tests/test_qwen3_4B_ppo_train_critic_only.py
index f0e1b19e08..64859a5716 100644
--- a/tests/test_qwen3_4B_ppo_train_critic_only.py
+++ b/tests/test_qwen3_4B_ppo_train_critic_only.py
@@ -1,5 +1,4 @@
 import os
-import tempfile
 
 import slime.utils.external_utils.command_utils as U
 
@@ -22,22 +21,6 @@ def prepare():
 
 
 def execute():
-    megatron_config = tempfile.NamedTemporaryFile("w", suffix=".yaml", delete=False)
-    megatron_config.write(
-        """
-megatron:
-  - name: default
-    role: critic
-    overrides:
-      lr: 1e-5
-  - name: default
-    role: actor
-    overrides:
-      lr: 1e-6
-"""
-    )
-    megatron_config.close()
-
     ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/{MODEL_NAME}_torch_dist "
 
     rollout_args = (
@@ -84,8 +67,9 @@ def execute():
         "--kl-coef 0.00 "
         "--entropy-coef 0.00 "
         "--eps-clip 4e-4 "
-        "--num-critic-only-steps 3 "
+        "--critic-train-only "
         "--normalize-advantages "
+        "--critic-lr 1e-5 "
     )
 
     optimizer_args = (
@@ -99,7 +83,7 @@ def execute():
 
     sglang_args = (
         "--rollout-num-gpus-per-engine 2 "
-        "--rollout-num-gpus 8 "
+        "--rollout-num-gpus 4 "
         "--sglang-mem-fraction-static 0.8 "
         "--sglang-max-running-requests 512 "
         "--sglang-enable-metrics "
@@ -116,13 +100,13 @@ def execute():
         "--attention-softmax-in-fp32 "
         # need to comment this when using model with MLA
         "--attention-backend flash "
-        "--actor-num-nodes 1 "
-        "--actor-num-gpus-per-node 8 "
-        "--colocate "
+        "--actor-num-nodes 0 "
+        "--actor-num-gpus-per-node 0 "
+        "--critic-num-nodes 1 "
+        "--critic-num-gpus-per-node 4 "
     )
 
     train_args = (
-        f"--megatron-config-path {megatron_config.name} "
         f"{ckpt_args} "
         f"{rollout_args} "
         f"{optimizer_args} "
diff --git a/tests/test_qwen3_5_mtp_bridge_mapping.py b/tests/test_qwen3_5_mtp_bridge_mapping.py
deleted file mode 100644
index cb304ab5db..0000000000
--- a/tests/test_qwen3_5_mtp_bridge_mapping.py
+++ /dev/null
@@ -1,242 +0,0 @@
-import importlib.util
-import sys
-import types
-from pathlib import Path
-
-import pytest
-import torch
-
-
-def install_bridge_stubs():
-    megatron_mod = types.ModuleType("megatron")
-    core_mod = types.ModuleType("megatron.core")
-    models_mod = types.ModuleType("megatron.core.models")
-    gpt_mod = types.ModuleType("megatron.core.models.gpt")
-    gpt_layer_specs_mod = types.ModuleType("megatron.core.models.gpt.gpt_layer_specs")
-    gpt_layer_specs_mod.get_gpt_mtp_block_spec = lambda _config, transformer_layer_spec, **_kwargs: (
-        "mtp-spec",
-        transformer_layer_spec,
-    )
-
-    mbridge_mod = types.ModuleType("mbridge")
-    mbridge_core_mod = types.ModuleType("mbridge.core")
-    mbridge_models_mod = types.ModuleType("mbridge.models")
-
-    def register_model(_names):
-        def decorator(cls):
-            return cls
-
-        return decorator
-
-    class Qwen2MoEBridge:
-        _MLP_MAPPING = {
-            "shared_experts.linear_fc1.weight": [
-                "model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight",
-                "model.layers.{layer_number}.mlp.shared_expert.up_proj.weight",
-            ],
-            "pre_mlp_layernorm": ["model.layers.{layer_number}.post_attention_layernorm.weight"],
-            "shared_experts.linear_fc2.weight": ["model.layers.{layer_number}.mlp.shared_expert.down_proj.weight"],
-            "mlp.router.weight": ["model.layers.{layer_number}.mlp.gate.weight"],
-            "shared_experts.gate_weight": ["model.layers.{layer_number}.mlp.shared_expert_gate.weight"],
-            "mlp.experts.linear_fc1": [
-                "model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight",
-                "model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight",
-            ],
-            "mlp.experts.linear_fc2": ["model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight"],
-        }
-
-        def _weight_name_mapping_mlp(self, name: str) -> list[str]:
-            layer_number = name.split(".")[2]
-            convert_names = []
-            for keyword, mapping_names in self._MLP_MAPPING.items():
-                if keyword in name:
-                    if "{expert_id}" in mapping_names[0]:
-                        expert_id = name.split("weight")[-1]
-                        convert_names.extend(
-                            [x.format(layer_number=layer_number, expert_id=expert_id) for x in mapping_names]
-                        )
-                    else:
-                        convert_names.extend([x.format(layer_number=layer_number) for x in mapping_names])
-                    break
-            if len(convert_names) == 0:
-                raise NotImplementedError(f"Unsupported parameter name: {name}")
-            return convert_names
-
-        def _weight_name_mapping_attention(self, name: str) -> list[str]:
-            raise NotImplementedError(f"Unexpected attention mapping lookup: {name}")
-
-        def _get_transformer_layer_spec(self, vp_stage=None):
-            return "REAL_LAYER_SPEC" if vp_stage is None else f"REAL_LAYER_SPEC_VP{vp_stage}"
-
-        def _get_gptmodel_args(self) -> dict:
-            return {"base": "ok"}
-
-        def _model_provider(self, callbacks):
-            def provider(pre_process, post_process, vp_stage=None):
-                transformer_layer_spec = self._get_transformer_layer_spec(vp_stage)
-                gptmodel_args = self._get_gptmodel_args()
-                return {"transformer_layer_spec": transformer_layer_spec, **gptmodel_args}
-
-            return provider
-
-        def _weight_to_mcore_format(self, _mcore_weights_name, hf_weights):
-            assert len(hf_weights) == 1
-            return hf_weights[0]
-
-        def _weight_to_hf_format(self, mcore_weights_name, mcore_weights):
-            return [mcore_weights_name], [mcore_weights]
-
-        def _build_base_config(self, **kwargs):
-            return kwargs
-
-    mbridge_core_mod.register_model = register_model
-    mbridge_models_mod.Qwen2MoEBridge = Qwen2MoEBridge
-
-    sys.modules["megatron"] = megatron_mod
-    sys.modules["megatron.core"] = core_mod
-    sys.modules["megatron.core.models"] = models_mod
-    sys.modules["megatron.core.models.gpt"] = gpt_mod
-    sys.modules["megatron.core.models.gpt.gpt_layer_specs"] = gpt_layer_specs_mod
-    sys.modules["mbridge"] = mbridge_mod
-    sys.modules["mbridge.core"] = mbridge_core_mod
-    sys.modules["mbridge.models"] = mbridge_models_mod
-
-
-def load_bridge_module():
-    install_bridge_stubs()
-    module_path = Path(__file__).resolve().parents[1] / "slime_plugins" / "mbridge" / "qwen3_5.py"
-    module_name = "test_qwen3_5_bridge_module"
-    sys.modules.pop(module_name, None)
-    spec = importlib.util.spec_from_file_location(module_name, module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
-def load_raw_export_module():
-    module_path = (
-        Path(__file__).resolve().parents[1] / "slime" / "backends" / "megatron_utils" / "megatron_to_hf" / "qwen3_5.py"
-    )
-    module_name = "test_qwen3_5_raw_export_module"
-    sys.modules.pop(module_name, None)
-    spec = importlib.util.spec_from_file_location(module_name, module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec.loader is not None
-    spec.loader.exec_module(module)
-    return module
-
-
-@pytest.mark.unit
-def test_mtp_moe_expert_mapping_uses_individual_hf_weights():
-    module = load_bridge_module()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-
-    fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc1.weight42")
-    fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.experts.linear_fc2.weight42")
-
-    assert fc1_names == [
-        "mtp.layers.0.mlp.experts.42.gate_proj.weight",
-        "mtp.layers.0.mlp.experts.42.up_proj.weight",
-    ]
-    assert fc2_names == ["mtp.layers.0.mlp.experts.42.down_proj.weight"]
-
-
-@pytest.mark.unit
-def test_mtp_dense_mlp_mapping_still_uses_dense_hf_weights():
-    module = load_bridge_module()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-
-    fc1_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc1.weight")
-    fc2_names = bridge._convert_mtp_param("mtp.layers.0.transformer_layer.mlp.linear_fc2.weight")
-
-    assert fc1_names == ["mtp.layers.0.mlp.gate_proj.weight", "mtp.layers.0.mlp.up_proj.weight"]
-    assert fc2_names == ["mtp.layers.0.mlp.down_proj.weight"]
-
-
-@pytest.mark.unit
-def test_mtp_block_spec_uses_current_transformer_layer_spec():
-    module = load_bridge_module()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-    bridge.config = "CONFIG_OBJECT"
-    bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1))
-
-    provider = bridge._model_provider([])
-    result = provider(True, True, vp_stage=3)
-
-    assert result["transformer_layer_spec"] == "REAL_LAYER_SPEC_VP3"
-    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()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-
-    weight = torch.arange(24, dtype=torch.float32).view(3, 8)
-    converted = bridge._weight_to_mcore_format("mtp.layers.0.eh_proj.weight", [weight])
-
-    assert torch.equal(converted, weight)
-
-
-@pytest.mark.unit
-def test_build_config_enables_gated_attention_when_transformer_config_supports_it():
-    module = load_bridge_module()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-    bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1))
-    bridge.TransformerConfigClass = types.SimpleNamespace(
-        __dataclass_fields__={
-            "mtp_num_layers": None,
-            "attention_output_gate": None,
-            "use_gated_attention": None,
-        }
-    )
-
-    config = bridge._build_config()
-
-    assert config["mtp_num_layers"] == 1
-    assert config["attention_output_gate"] is True
-    assert config["use_gated_attention"] is True
-
-
-@pytest.mark.unit
-def test_build_config_skips_gated_attention_when_transformer_config_does_not_support_it():
-    module = load_bridge_module()
-    bridge = module.Qwen3_5Bridge.__new__(module.Qwen3_5Bridge)
-    bridge.hf_config = types.SimpleNamespace(text_config=types.SimpleNamespace(mtp_num_hidden_layers=1))
-    bridge.TransformerConfigClass = types.SimpleNamespace(
-        __dataclass_fields__={
-            "mtp_num_layers": None,
-            "attention_output_gate": None,
-        }
-    )
-
-    config = bridge._build_config()
-
-    assert config["mtp_num_layers"] == 1
-    assert config["attention_output_gate"] is True
-    assert "use_gated_attention" not in config
-
-
-@pytest.mark.unit
-def test_raw_qwen3_5_mtp_export_keeps_eh_proj_column_order():
-    module = load_raw_export_module()
-
-    weight = torch.arange(24, dtype=torch.float32).view(3, 8)
-    converted = module.convert_qwen3_5_to_hf(
-        types.SimpleNamespace(), "module.module.mtp.layers.0.eh_proj.weight", weight
-    )
-
-    assert converted == [("mtp.fc.weight", weight)]
diff --git a/tests/test_qwen3_linear_attention_cu_seqlens.py b/tests/test_qwen3_linear_attention_cu_seqlens.py
deleted file mode 100644
index 7107561fa5..0000000000
--- a/tests/test_qwen3_linear_attention_cu_seqlens.py
+++ /dev/null
@@ -1,165 +0,0 @@
-from __future__ import annotations
-
-import importlib
-import sys
-import types
-from types import SimpleNamespace
-
-import pytest
-import torch
-import torch.nn as nn
-
-
-def install_megatron_stubs() -> None:
-    if "megatron" in sys.modules:
-        return
-
-    megatron_mod = types.ModuleType("megatron")
-    core_mod = types.ModuleType("megatron.core")
-    models_mod = types.ModuleType("megatron.core.models")
-    gpt_mod = types.ModuleType("megatron.core.models.gpt")
-    gpt_layer_specs_mod = types.ModuleType("megatron.core.models.gpt.gpt_layer_specs")
-    inference_mod = types.ModuleType("megatron.core.inference")
-    inference_contexts_mod = types.ModuleType("megatron.core.inference.contexts")
-    packed_seq_mod = types.ModuleType("megatron.core.packed_seq_params")
-    transformer_mod = types.ModuleType("megatron.core.transformer")
-    transformer_module_mod = types.ModuleType("megatron.core.transformer.module")
-    spec_utils_mod = types.ModuleType("megatron.core.transformer.spec_utils")
-    transformer_block_mod = types.ModuleType("megatron.core.transformer.transformer_block")
-    transformer_layer_mod = types.ModuleType("megatron.core.transformer.transformer_layer")
-
-    class PackedSeqParams:
-        def __init__(self, **kwargs):
-            for key, value in kwargs.items():
-                setattr(self, key, value)
-
-    class MegatronModule(nn.Module):
-        def __init__(self, config=None):
-            super().__init__()
-            self.config = config
-
-    class ModuleSpec:
-        def __init__(self, module=None, params=None):
-            self.module = module
-            self.params = params or {}
-
-    mpu_stub = types.SimpleNamespace(
-        get_context_parallel_world_size=lambda: 1,
-        get_context_parallel_group=lambda: None,
-        get_context_parallel_rank=lambda: 0,
-        get_tensor_model_parallel_group=lambda: None,
-    )
-    tensor_parallel_stub = types.SimpleNamespace(
-        gather_from_sequence_parallel_region=lambda x, group=None: x,
-        scatter_to_sequence_parallel_region=lambda x, group=None: x,
-    )
-
-    gpt_layer_specs_mod.get_gpt_decoder_block_spec = lambda *args, **kwargs: None
-    inference_contexts_mod.BaseInferenceContext = type("BaseInferenceContext", (), {})
-    packed_seq_mod.PackedSeqParams = PackedSeqParams
-    transformer_module_mod.MegatronModule = MegatronModule
-    spec_utils_mod.ModuleSpec = ModuleSpec
-    transformer_block_mod.get_num_layers_to_build = lambda *args, **kwargs: 0
-    transformer_layer_mod.get_transformer_layer_offset = lambda *args, **kwargs: 0
-
-    core_mod.mpu = mpu_stub
-    core_mod.tensor_parallel = tensor_parallel_stub
-
-    sys.modules["megatron"] = megatron_mod
-    sys.modules["megatron.core"] = core_mod
-    sys.modules["megatron.core.models"] = models_mod
-    sys.modules["megatron.core.models.gpt"] = gpt_mod
-    sys.modules["megatron.core.models.gpt.gpt_layer_specs"] = gpt_layer_specs_mod
-    sys.modules["megatron.core.inference"] = inference_mod
-    sys.modules["megatron.core.inference.contexts"] = inference_contexts_mod
-    sys.modules["megatron.core.packed_seq_params"] = packed_seq_mod
-    sys.modules["megatron.core.transformer"] = transformer_mod
-    sys.modules["megatron.core.transformer.module"] = transformer_module_mod
-    sys.modules["megatron.core.transformer.spec_utils"] = spec_utils_mod
-    sys.modules["megatron.core.transformer.transformer_block"] = transformer_block_mod
-    sys.modules["megatron.core.transformer.transformer_layer"] = transformer_layer_mod
-
-
-class FakeShortConvolution(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-
-    def forward(self, x, cu_seqlens=None, **kwargs):
-        return x, None
-
-
-class FakeFusedRMSNormGated(nn.Module):
-    def __init__(self, *args, **kwargs):
-        super().__init__()
-
-    def forward(self, x, z):
-        return x
-
-
-def make_config() -> SimpleNamespace:
-    return SimpleNamespace(
-        hidden_size=32,
-        linear_num_value_heads=4,
-        linear_num_key_heads=2,
-        linear_key_head_dim=8,
-        linear_value_head_dim=8,
-        linear_conv_kernel_dim=4,
-        hidden_act="silu",
-        rms_norm_eps=1e-6,
-        dtype=torch.float32,
-    )
-
-
-def load_module(module_name: str):
-    install_megatron_stubs()
-    sys.modules.pop("slime_plugins.models.hf_attention", None)
-    sys.modules.pop(module_name, None)
-    return importlib.import_module(module_name)
-
-
-@pytest.mark.unit
-@pytest.mark.parametrize(
-    ("module_name", "class_name"),
-    [
-        ("slime_plugins.models.qwen3_5", "Qwen3_5GatedDeltaNet"),
-        ("slime_plugins.models.qwen3_next", "Qwen3NextGatedDeltaNet"),
-    ],
-)
-def test_linear_attention_forwards_cu_seqlens_to_chunk_kernel(monkeypatch, module_name: str, class_name: str):
-    module = load_module(module_name)
-
-    monkeypatch.setattr(module.torch.cuda, "current_device", lambda: "cpu")
-    monkeypatch.setattr(module, "ShortConvolution", FakeShortConvolution)
-    monkeypatch.setattr(module, "FusedRMSNormGated", FakeFusedRMSNormGated)
-
-    chunk_calls = []
-
-    def fake_chunk_gated_delta_rule(
-        q,
-        k,
-        v,
-        *,
-        g,
-        beta,
-        initial_state,
-        output_final_state,
-        use_qk_l2norm_in_kernel,
-        cu_seqlens=None,
-        **kwargs,
-    ):
-        chunk_calls.append(cu_seqlens.clone() if cu_seqlens is not None else None)
-        assert q.shape[0] == 1
-        assert cu_seqlens is not None
-        return torch.zeros_like(v), None
-
-    monkeypatch.setattr(module, "chunk_gated_delta_rule", fake_chunk_gated_delta_rule)
-
-    layer = getattr(module, class_name)(make_config(), layer_idx=0)
-    hidden_states = torch.randn(1, 7, 32)
-    cu_seqlens = torch.tensor([0, 3, 7], dtype=torch.int32)
-
-    output = layer(hidden_states, cu_seqlens=cu_seqlens)
-
-    assert output.shape == hidden_states.shape
-    assert len(chunk_calls) == 1
-    assert torch.equal(chunk_calls[0], cu_seqlens)
diff --git a/tests/test_qwen3_vl_4B_fsdp.py b/tests/test_qwen3_vl_4B_fsdp.py
new file mode 100644
index 0000000000..5509cad29c
--- /dev/null
+++ b/tests/test_qwen3_vl_4B_fsdp.py
@@ -0,0 +1,117 @@
+import os
+import slime.utils.external_utils.command_utils as U
+
+ENABLE_EVAL = bool(int(os.environ.get("SLIME_TEST_ENABLE_EVAL", "1")))
+NUM_GPUS = 8
+
+MODEL_NAME = "Qwen3-VL-4B-Instruct"
+DATASET_NAME = "chenhegu/geo3k_imgurl"
+
+
+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(DATASET_NAME)
+
+
+def execute():
+    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME} "
+
+    rollout_args = (
+        "--prompt-data /root/datasets/geo3k_imgurl/train.parquet "
+        "--input-key problem "
+        "--label-key answer "
+        "--apply-chat-template "
+        "--rollout-shuffle "
+        "--rm-type math "
+        "--num-rollout 3 "
+        "--rollout-batch-size 8 "
+        "--n-samples-per-prompt 4 "
+        "--rollout-max-response-len 4096 "
+        "--rollout-temperature 1 "
+        "--global-batch-size 32 "
+    )
+
+    # multimodal keys required for vlm datasets
+    multimodal_args = '--multimodal-keys \'{"image": "images"}\' '
+
+    eval_args = (
+        f"{'--eval-interval 20 ' if ENABLE_EVAL else ''}"
+        "--eval-prompt-data geo3k /root/datasets/geo3k_imgurl/test.parquet "
+        "--n-samples-per-eval-prompt 1 "
+        "--eval-max-response-len 4096 "
+    )
+
+    fsdp_args = "--train-backend fsdp " "--gradient-checkpointing " "--update-weight-buffer-size 536870912 "
+
+    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 "
+    )
+
+    optimizer_args = (
+        "--optimizer adam "
+        "--lr 1e-6 "
+        "--lr-decay-style constant "
+        "--weight-decay 0.1 "
+        "--adam-beta1 0.9 "
+        "--adam-beta2 0.98 "
+    )
+
+    sglang_args = (
+        "--rollout-num-gpus-per-engine 1 "
+        "--sglang-mem-fraction-static 0.6 "
+        "--sglang-decode-log-interval 1000 "
+        "--sglang-enable-metrics "
+        # "--sglang-enable-deterministic-inference "
+        # "--sglang-rl-on-policy-target fsdp "
+        "--sglang-attention-backend fa3 "
+        "--attn-implementation flash_attention_3 "
+        "--sglang-cuda-graph-max-bs 32 "
+        # "--deterministic-mode "
+        # "--true-on-policy-mode "
+    )
+
+    ci_args = "--ci-test "
+
+    misc_args = "--actor-num-nodes 1 " f"--actor-num-gpus-per-node {NUM_GPUS} " "--colocate "
+
+    train_args = (
+        f"{ckpt_args} "
+        f"{rollout_args} "
+        f"{multimodal_args} "
+        f"{optimizer_args} "
+        f"{grpo_args} "
+        f"{U.get_default_wandb_args(__file__)} "
+        f"{fsdp_args} "
+        f"{eval_args} "
+        f"{sglang_args} "
+        f"{ci_args} "
+        f"{misc_args} "
+    )
+
+    extra_env_vars = {
+        # "NCCL_ALGO": "allreduce:tree",
+        # "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0",
+        # "CUBLAS_WORKSPACE_CONFIG": ":4096:8",
+        "CUDA_DEVICE_MAX_CONNECTIONS": "1",
+    }
+
+    U.execute_train(
+        train_args=train_args,
+        num_gpus_per_node=NUM_GPUS,
+        megatron_model_type=None,
+        extra_env_vars=extra_env_vars,
+    )
+
+
+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_sglang_config_mixed_offload.py b/tests/test_sglang_config_mixed_offload.py
deleted file mode 100644
index 90d3d97389..0000000000
--- a/tests/test_sglang_config_mixed_offload.py
+++ /dev/null
@@ -1,158 +0,0 @@
-"""E2E test: mixed offload with updatable + frozen models.
-
-Deploys two models via --sglang-config in colocate mode:
-  - "actor": update_weights=true, 4 GPUs → overlaps with megatron, gets offloaded
-    and weights updated from training.
-  - "ref":   update_weights=false, 4 GPUs → overlaps with megatron, gets offloaded
-    and weights restored from disk (update_weights_from_disk).
-
-Key coverage:
-  - Per-group needs_offload (both overlap with megatron in colocate mode)
-  - update_weights_from_disk for frozen model
-  - Selective flush_cache (only for offloaded / updatable engines)
-  - Offload/onload cycle completes without crash
-"""
-
-import os
-import tempfile
-
-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"
-NUM_GPUS = 8
-
-# Two models on 8 GPUs (colocate): actor gets weight updates, ref is frozen.
-SGLANG_CONFIG_YAML = """\
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 1
-  - name: ref
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 1
-"""
-
-
-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/gsm8k")
-
-
-def execute():
-    config_file = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", prefix="sglang_mixed_offload_", delete=False)
-    config_file.write(SGLANG_CONFIG_YAML)
-    config_file.flush()
-    config_path = config_file.name
-
-    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
-
-    rollout_args = (
-        "--prompt-data /root/datasets/gsm8k/train.parquet "
-        "--input-key messages "
-        "--label-key label "
-        "--apply-chat-template "
-        "--rollout-shuffle "
-        "--rm-type math "
-        "--num-rollout 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 512 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-    )
-
-    eval_args = (
-        "--eval-interval 20 "
-        "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet "
-        "--n-samples-per-eval-prompt 1 "
-        "--eval-max-response-len 512 "
-        "--eval-top-k 1 "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 1 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 4096 "
-    )
-
-    grpo_args = (
-        "--advantage-estimator grpo "
-        "--use-kl-loss "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type low_var_kl "
-        "--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 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 1 "
-        f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
-        "--sglang-cuda-graph-max-bs 32 "
-        f"--sglang-config {config_path} "
-    )
-
-    ci_args = "--ci-test "
-
-    misc_args = (
-        "--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 "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    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"{ci_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()
-    os.environ.pop("http_proxy", None)
-    os.environ.pop("https_proxy", None)
-    os.environ.pop("HTTP_PROXY", None)
-    os.environ.pop("HTTPS_PROXY", None)
-    execute()
diff --git a/tests/test_sglang_config_mixed_offload_ft.py b/tests/test_sglang_config_mixed_offload_ft.py
deleted file mode 100644
index f017965c5f..0000000000
--- a/tests/test_sglang_config_mixed_offload_ft.py
+++ /dev/null
@@ -1,165 +0,0 @@
-"""E2E test: mixed offload with fault tolerance.
-
-Same two-model layout as test_sglang_config_mixed_offload.py but with
-fault tolerance enabled.  --ci-test triggers a simulated engine crash
-on the updatable (actor) server, testing:
-  - Health monitor detects crash and marks engine as None
-  - RolloutServer.recover() restarts the dead engine
-  - Updatable engines: offload → resume_memory_occupation → update_weights
-  - Non-updatable engines: offload → update_weights_from_disk
-  - Training continues after recovery
-"""
-
-import os
-import tempfile
-
-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"
-NUM_GPUS = 8
-
-# Two models on 8 GPUs (colocate): actor gets weight updates, ref is frozen.
-SGLANG_CONFIG_YAML = """\
-sglang:
-  - name: actor
-    update_weights: true
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 1
-  - name: ref
-    update_weights: false
-    server_groups:
-      - worker_type: regular
-        num_gpus: 4
-        num_gpus_per_engine: 1
-"""
-
-
-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/gsm8k")
-
-
-def execute():
-    config_file = tempfile.NamedTemporaryFile(
-        mode="w", suffix=".yaml", prefix="sglang_mixed_offload_ft_", delete=False
-    )
-    config_file.write(SGLANG_CONFIG_YAML)
-    config_file.flush()
-    config_path = config_file.name
-
-    ckpt_args = f"--hf-checkpoint /root/models/{MODEL_NAME}/ " f"--ref-load /root/models/{MODEL_NAME}/ "
-
-    rollout_args = (
-        "--prompt-data /root/datasets/gsm8k/train.parquet "
-        "--input-key messages "
-        "--label-key label "
-        "--apply-chat-template "
-        "--rollout-shuffle "
-        "--rm-type math "
-        "--num-rollout 3 "
-        "--rollout-batch-size 8 "
-        "--n-samples-per-prompt 4 "
-        "--rollout-max-response-len 512 "
-        "--rollout-temperature 0.8 "
-        "--global-batch-size 32 "
-    )
-
-    eval_args = (
-        "--eval-interval 20 "
-        "--eval-prompt-data gsm8k /root/datasets/gsm8k/test.parquet "
-        "--n-samples-per-eval-prompt 1 "
-        "--eval-max-response-len 512 "
-        "--eval-top-k 1 "
-    )
-
-    perf_args = (
-        "--tensor-model-parallel-size 1 "
-        "--sequence-parallel "
-        "--pipeline-model-parallel-size 1 "
-        "--context-parallel-size 1 "
-        "--expert-model-parallel-size 1 "
-        "--expert-tensor-parallel-size 1 "
-        "--use-dynamic-batch-size "
-        "--max-tokens-per-gpu 4096 "
-    )
-
-    grpo_args = (
-        "--advantage-estimator grpo "
-        "--use-kl-loss "
-        "--kl-loss-coef 0.00 "
-        "--kl-loss-type low_var_kl "
-        "--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 "
-    )
-
-    sglang_args = (
-        "--rollout-num-gpus-per-engine 1 "
-        f"--sglang-mem-fraction-static {0.6 if TIGHT_DEVICE_MEMORY else 0.7} "
-        "--sglang-cuda-graph-max-bs 32 "
-        f"--sglang-config {config_path} "
-    )
-
-    ci_args = "--ci-test "
-
-    fault_tolerance_args = (
-        "--use-fault-tolerance "
-        "--rollout-health-check-interval 5 "
-        "--rollout-health-check-timeout 10 "
-        "--rollout-health-check-first-wait 0 "
-    )
-
-    misc_args = (
-        "--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 "
-        "--megatron-to-hf-mode bridge "
-    )
-
-    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"{ci_args} "
-        f"{fault_tolerance_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()
-    os.environ.pop("http_proxy", None)
-    os.environ.pop("https_proxy", None)
-    os.environ.pop("HTTP_PROXY", None)
-    os.environ.pop("HTTPS_PROXY", None)
-    execute()
diff --git a/tests/utils/test_loss_mask_type_qwen35.py b/tests/utils/test_loss_mask_type_qwen35.py
deleted file mode 100644
index 69ffd1513f..0000000000
--- a/tests/utils/test_loss_mask_type_qwen35.py
+++ /dev/null
@@ -1,236 +0,0 @@
-from slime.utils.mask_utils import MultiTurnLossMaskGenerator
-
-
-class FakeQwen35Tokenizer:
-    """A tiny char-level tokenizer that models the Qwen3.5 assistant formatting rule.
-
-    The critical behavior we need for these tests is:
-    1. `add_generation_prompt=True` appends `<|im_start|>assistant\n<think>\n`
-    2. Only assistant turns after the last non-tool user query are wrapped in
-       `<think>...</think>`
-    """
-
-    assistant_generation_prompt = "<|im_start|>assistant\n<think>\n"
-
-    def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False):
-        encoded = {"input_ids": [ord(ch) for ch in text]}
-        if return_offsets_mapping:
-            encoded["offset_mapping"] = [(index, index + 1) for index in range(len(text))]
-        return encoded
-
-    def decode(self, token_ids):
-        return "".join(chr(token_id) for token_id in token_ids)
-
-    def apply_chat_template(
-        self,
-        messages,
-        tokenize=True,
-        tools=None,
-        add_generation_prompt=False,
-        return_dict=False,
-        add_special_tokens=False,
-        **kwargs,
-    ):
-        rendered = self.render(messages, add_generation_prompt=add_generation_prompt)
-        if tokenize:
-            return [ord(ch) for ch in rendered]
-        return rendered
-
-    def render(self, messages, add_generation_prompt=False):
-        rendered, _ = self.render_with_expected_mask(messages, add_generation_prompt=add_generation_prompt)
-        return rendered
-
-    def render_with_expected_mask(self, messages, add_generation_prompt=False):
-        if not messages:
-            raise ValueError("No messages provided.")
-
-        pieces = []
-        mask = []
-        last_query_index = self._find_last_query_index(messages)
-        has_assistant = any(message["role"] == "assistant" for message in messages)
-        if has_assistant and last_query_index is None:
-            raise ValueError("No user query found in messages.")
-
-        for index, message in enumerate(messages):
-            role = message["role"]
-            if role == "system":
-                if index != 0:
-                    raise ValueError("System message must be at the beginning.")
-                piece = f"<|im_start|>system\n{message['content']}<|im_end|>\n"
-                pieces.append(piece)
-                mask.extend([0] * len(piece))
-                continue
-
-            if role == "user":
-                piece = f"<|im_start|>user\n{message['content']}<|im_end|>\n"
-                pieces.append(piece)
-                mask.extend([0] * len(piece))
-                continue
-
-            if role == "tool":
-                piece = ""
-                if index > 0 and messages[index - 1]["role"] != "tool":
-                    piece += "<|im_start|>user"
-                piece += f"\n<tool_response>\n{message['content']}\n</tool_response>"
-                if index == len(messages) - 1 or messages[index + 1]["role"] != "tool":
-                    piece += "<|im_end|>\n"
-                pieces.append(piece)
-                mask.extend([0] * len(piece))
-                continue
-
-            if role != "assistant":
-                raise NotImplementedError(f"Unsupported role in test tokenizer: {role}")
-
-            reasoning, answer = self._split_assistant_content(message["content"])
-            if index > last_query_index:
-                prefix = "<|im_start|>assistant\n<think>\n"
-                target = f"{reasoning}\n</think>\n\n{answer}"
-            else:
-                prefix = "<|im_start|>assistant\n"
-                target = answer
-
-            target += self._render_tool_calls(answer, message.get("tool_calls"))
-            target += "<|im_end|>\n"
-
-            pieces.append(prefix + target)
-            mask.extend([0] * len(prefix))
-            mask.extend([1] * len(target))
-
-        if add_generation_prompt:
-            pieces.append(self.assistant_generation_prompt)
-            mask.extend([0] * len(self.assistant_generation_prompt))
-
-        return "".join(pieces), mask
-
-    @staticmethod
-    def _split_assistant_content(content):
-        if "</think>" not in content:
-            return "", content
-        reasoning = content.split("</think>")[0].split("<think>")[-1].strip("\n")
-        answer = content.split("</think>")[-1].lstrip("\n")
-        return reasoning, answer
-
-    @staticmethod
-    def _render_tool_calls(answer, tool_calls):
-        if not tool_calls:
-            return ""
-
-        pieces = []
-        for index, tool_call in enumerate(tool_calls):
-            function_call = tool_call.get("function", tool_call)
-            function_name = function_call["name"]
-            if index == 0:
-                if answer.strip():
-                    pieces.append(f"\n\n<tool_call>\n<function={function_name}>\n")
-                else:
-                    pieces.append(f"<tool_call>\n<function={function_name}>\n")
-            else:
-                pieces.append(f"\n<tool_call>\n<function={function_name}>\n")
-
-            for argument_name, argument_value in function_call.get("arguments", {}).items():
-                pieces.append(f"<parameter={argument_name}>\n")
-                pieces.append(str(argument_value))
-                pieces.append("\n</parameter>\n")
-
-            pieces.append("</function>\n</tool_call>")
-
-        return "".join(pieces)
-
-    @staticmethod
-    def _find_last_query_index(messages):
-        last_query_index = None
-        for index, message in enumerate(messages):
-            if message["role"] == "user":
-                last_query_index = index
-        return last_query_index
-
-
-def test_qwen3_and_qwen3_5_match_on_single_turn_qwen35_data():
-    tokenizer = FakeQwen35Tokenizer()
-    messages = [
-        {"role": "system", "content": "SYSTEM"},
-        {"role": "user", "content": "USER"},
-        {"role": "assistant", "content": "<think>REASONING</think>\nANSWER"},
-    ]
-
-    expected_text, expected_mask = tokenizer.render_with_expected_mask(messages)
-    expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"]
-
-    for loss_mask_type in ("qwen3", "qwen3_5"):
-        generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type=loss_mask_type)
-        token_ids, loss_mask = generator.get_loss_mask(messages)
-        assert token_ids == expected_token_ids
-        assert loss_mask == expected_mask
-
-
-def test_qwen3_and_qwen3_5_diverge_on_multi_turn_qwen35_data():
-    tokenizer = FakeQwen35Tokenizer()
-    messages = [
-        {"role": "system", "content": "SYSTEM"},
-        {"role": "user", "content": "USER_1"},
-        {"role": "assistant", "content": "ANSWER_1"},
-        {"role": "user", "content": "USER_2"},
-        {"role": "assistant", "content": "<think>REASONING_2</think>\nANSWER_2"},
-    ]
-
-    expected_text, expected_mask = tokenizer.render_with_expected_mask(messages)
-    expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"]
-
-    qwen3_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3")
-    qwen35_generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3_5")
-
-    qwen3_token_ids, qwen3_loss_mask = qwen3_generator.get_loss_mask(messages)
-    qwen35_token_ids, qwen35_loss_mask = qwen35_generator.get_loss_mask(messages)
-
-    assert qwen3_token_ids != qwen35_token_ids
-    assert qwen3_loss_mask != qwen35_loss_mask
-
-    # `qwen3` is incompatible with the real Qwen3.5 full-text rendering:
-    # it rebuilds each assistant turn in isolation and fabricates a `<think>` block
-    # for the earlier assistant answer.
-    assert qwen3_token_ids != expected_token_ids
-    assert qwen3_loss_mask != expected_mask
-
-    # `qwen3_5` matches the expected full-text rendering and supervises every
-    # assistant turn in the full conversation.
-    assert qwen35_token_ids == expected_token_ids
-    assert qwen35_loss_mask == expected_mask
-
-    assert qwen3_generator.get_text_from_loss_mask(qwen3_token_ids, qwen3_loss_mask) == [
-        "\n</think>\n\nANSWER_1<|im_end|>\n",
-        "REASONING_2\n</think>\n\nANSWER_2<|im_end|>\n",
-    ]
-    expected_selected_texts = [
-        "ANSWER_1<|im_end|>\n",
-        "REASONING_2\n</think>\n\nANSWER_2<|im_end|>\n",
-    ]
-    assert qwen35_generator.get_text_from_loss_mask(qwen35_token_ids, qwen35_loss_mask) == expected_selected_texts
-    assert qwen35_generator.get_text_from_loss_mask(expected_token_ids, expected_mask) == expected_selected_texts
-
-
-def test_qwen3_5_matches_expected_mask_for_tool_call_flow():
-    tokenizer = FakeQwen35Tokenizer()
-    messages = [
-        {"role": "system", "content": "SYSTEM"},
-        {"role": "user", "content": "USER"},
-        {
-            "role": "assistant",
-            "content": "TOOL_CALL",
-            "tool_calls": [{"function": {"name": "terminal", "arguments": {"command": "ls"}}}],
-        },
-        {"role": "tool", "content": "README.md"},
-        {"role": "assistant", "content": "<think>REASONING</think>\nFINAL"},
-    ]
-
-    expected_text, expected_mask = tokenizer.render_with_expected_mask(messages)
-    expected_token_ids = tokenizer(expected_text, add_special_tokens=False)["input_ids"]
-
-    generator = MultiTurnLossMaskGenerator(tokenizer, tokenizer_type="qwen3_5")
-    token_ids, loss_mask = generator.get_loss_mask(messages)
-
-    assert token_ids == expected_token_ids
-    assert loss_mask == expected_mask
-    assert generator.get_text_from_loss_mask(token_ids, loss_mask) == [
-        "\n</think>\n\nTOOL_CALL\n\n<tool_call>\n<function=terminal>\n<parameter=command>\nls\n</parameter>\n</function>\n</tool_call><|im_end|>\n",
-        "REASONING\n</think>\n\nFINAL<|im_end|>\n",
-    ]
diff --git a/tests/utils/test_megatron_bridge_utils.py b/tests/utils/test_megatron_bridge_utils.py
deleted file mode 100644
index d47babcaca..0000000000
--- a/tests/utils/test_megatron_bridge_utils.py
+++ /dev/null
@@ -1,64 +0,0 @@
-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/tests/utils/test_megatron_role_config.py b/tests/utils/test_megatron_role_config.py
deleted file mode 100644
index 24cc37c57f..0000000000
--- a/tests/utils/test_megatron_role_config.py
+++ /dev/null
@@ -1,158 +0,0 @@
-"""Unit tests for Megatron role config parsing and application."""
-
-import tempfile
-from argparse import Namespace
-
-import pytest
-import yaml
-
-
-def _write_yaml(data: dict) -> str:
-    handle = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
-    yaml.dump(data, handle)
-    handle.flush()
-    return handle.name
-
-
-def _base_args(**overrides):
-    args = dict(
-        lr=2e-6,
-        tensor_model_parallel_size=1,
-        kl_coef=0.1,
-        use_kl_loss=False,
-        use_opd=True,
-        opd_type="megatron",
-        custom_advantage_function_path="slime.test.adv",
-        untie_embeddings_and_output_weights=False,
-        actor_num_nodes=1,
-        actor_num_gpus_per_node=1,
-        critic_num_nodes=1,
-        critic_num_gpus_per_node=1,
-        use_critic=False,
-        megatron_config_path=None,
-        start_rollout_id=None,
-        rollout_global_dataset=False,
-    )
-    args.update(overrides)
-    return Namespace(**args)
-
-
-class TestMegatronRoleConfig:
-    def test_parse_actor_and_critic_role_overrides(self):
-        from slime.utils.arguments import parse_megatron_role_args
-
-        path = _write_yaml(
-            {
-                "megatron": [
-                    {
-                        "name": "default",
-                        "role": "critic",
-                        "overrides": {"lr": "1e-5", "tensor_model_parallel_size": 2},
-                    },
-                    {"name": "default", "role": "actor", "overrides": {"lr": "1e-6", "tensor_model_parallel_size": 4}},
-                ]
-            }
-        )
-        args = _base_args()
-
-        actor_args = parse_megatron_role_args(args, path, role="actor")
-        critic_args = parse_megatron_role_args(args, path, role="critic")
-
-        assert actor_args.lr == 1e-6
-        assert actor_args.tensor_model_parallel_size == 4
-        assert actor_args.kl_coef == args.kl_coef
-        assert actor_args.use_opd is args.use_opd
-
-        assert critic_args.lr == 1e-5
-        assert critic_args.tensor_model_parallel_size == 2
-        assert critic_args.kl_coef == 0
-        assert critic_args.use_opd is False
-        assert critic_args.custom_advantage_function_path is None
-        assert critic_args.untie_embeddings_and_output_weights is True
-
-    def test_missing_role_inherits_base_args(self):
-        from slime.utils.arguments import parse_megatron_role_args
-
-        path = _write_yaml(
-            {
-                "megatron": [
-                    {"name": "default", "role": "actor", "overrides": {"lr": "1e-6"}},
-                ]
-            }
-        )
-        args = _base_args()
-
-        critic_args = parse_megatron_role_args(args, path, role="critic")
-
-        assert critic_args is not args
-        assert critic_args.lr == args.lr
-        assert critic_args.kl_coef == 0
-        assert critic_args.use_opd is False
-
-    @pytest.mark.parametrize(
-        "config",
-        [
-            {"critic": [{"name": "default", "overrides": {"lr": "1e-5"}}]},
-            {"lr": "1e-5"},
-            {},
-        ],
-    )
-    def test_requires_top_level_megatron_key(self, config):
-        from slime.utils.arguments import parse_megatron_role_args
-
-        path = _write_yaml(config)
-        args = _base_args()
-
-        with pytest.raises(AssertionError, match="top-level 'megatron' list"):
-            parse_megatron_role_args(args, path, role="critic")
-
-    def test_create_training_models_applies_actor_override_without_critic(self, monkeypatch):
-        from slime.ray import placement_group as placement_group_module
-
-        path = _write_yaml(
-            {
-                "megatron": [
-                    {"name": "default", "role": "actor", "overrides": {"lr": "1e-6"}},
-                ]
-            }
-        )
-        args = _base_args(megatron_config_path=path, use_critic=False)
-
-        class DummyModel:
-            def __init__(self, model_args):
-                self.args = model_args
-                self.init_calls = []
-                self.rollout_manager = None
-
-            def async_init(self, model_args, role, with_ref=False, with_opd_teacher=False):
-                self.args = model_args
-                self.init_calls.append(
-                    {
-                        "args": model_args,
-                        "role": role,
-                        "with_ref": with_ref,
-                        "with_opd_teacher": with_opd_teacher,
-                    }
-                )
-                return [7]
-
-            def set_rollout_manager(self, rollout_manager):
-                self.rollout_manager = rollout_manager
-
-        def fake_allocate_train_group(args, num_nodes, num_gpus_per_node, pg, role="actor"):
-            return DummyModel(args)
-
-        monkeypatch.setattr(placement_group_module, "allocate_train_group", fake_allocate_train_group)
-        monkeypatch.setattr(placement_group_module.ray, "get", lambda value: value)
-
-        actor_model, critic_model = placement_group_module.create_training_models(
-            args,
-            {"actor": None, "critic": None},
-            object(),
-        )
-
-        assert critic_model is None
-        assert actor_model.args.lr == 1e-6
-        assert actor_model.init_calls[0]["args"].lr == 1e-6
-        assert actor_model.init_calls[0]["role"] == "actor"
-        assert args.start_rollout_id == 7
diff --git a/tests/utils/test_sglang_config.py b/tests/utils/test_sglang_config.py
deleted file mode 100644
index 5015a74558..0000000000
--- a/tests/utils/test_sglang_config.py
+++ /dev/null
@@ -1,133 +0,0 @@
-"""Unit tests for SglangConfig multi-model parsing with update_weights."""
-
-import tempfile
-
-import pytest
-import yaml
-
-
-def _write_yaml(data: dict) -> str:
-    f = tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False)
-    yaml.dump(data, f)
-    f.flush()
-    return f.name
-
-
-class TestSglangConfigUpdateWeights:
-    def test_update_weights_default_true(self):
-        """Models without explicit update_weights should default to True."""
-        from slime.backends.sglang_utils.sglang_config import SglangConfig
-
-        path = _write_yaml(
-            {
-                "sglang": [
-                    {
-                        "name": "actor",
-                        "engine_groups": [{"worker_type": "regular", "num_gpus": 4}],
-                    }
-                ]
-            }
-        )
-        config = SglangConfig.from_yaml(path)
-        assert len(config.models) == 1
-        assert config.models[0].update_weights is True
-
-    def test_update_weights_explicit_false(self):
-        """Models with update_weights: false should be parsed correctly."""
-        from slime.backends.sglang_utils.sglang_config import SglangConfig
-
-        path = _write_yaml(
-            {
-                "sglang": [
-                    {
-                        "name": "actor",
-                        "update_weights": True,
-                        "engine_groups": [{"worker_type": "regular", "num_gpus": 4}],
-                    },
-                    {
-                        "name": "ref",
-                        "update_weights": False,
-                        "model_path": "/path/to/ref",
-                        "engine_groups": [{"worker_type": "regular", "num_gpus": 2}],
-                    },
-                ]
-            }
-        )
-        config = SglangConfig.from_yaml(path)
-        assert len(config.models) == 2
-        assert config.models[0].name == "actor"
-        assert config.models[0].update_weights is True
-        assert config.models[1].name == "ref"
-        assert config.models[1].update_weights is False
-        assert config.models[1].model_path == "/path/to/ref"
-
-    def test_multi_model_total_gpus(self):
-        """total_num_gpus should sum across all models."""
-        from slime.backends.sglang_utils.sglang_config import SglangConfig
-
-        path = _write_yaml(
-            {
-                "sglang": [
-                    {
-                        "name": "actor",
-                        "server_groups": [{"worker_type": "regular", "num_gpus": 8}],
-                    },
-                    {
-                        "name": "ref",
-                        "update_weights": False,
-                        "server_groups": [{"worker_type": "regular", "num_gpus": 4}],
-                    },
-                ]
-            }
-        )
-        config = SglangConfig.from_yaml(path)
-        assert config.total_num_gpus == 12
-
-
-class TestGetModelUrl:
-    def test_get_model_url_basic(self):
-        """get_model_url should return the correct URL for a named model."""
-        from argparse import Namespace
-
-        from slime.rollout.sglang_rollout import get_model_url
-
-        args = Namespace(
-            sglang_router_ip="10.0.0.1",
-            sglang_router_port=3000,
-            sglang_model_routers={
-                "actor": ("10.0.0.1", 3000),
-                "ref": ("10.0.0.1", 3001),
-            },
-        )
-        assert get_model_url(args, "actor") == "http://10.0.0.1:3000/generate"
-        assert get_model_url(args, "ref") == "http://10.0.0.1:3001/generate"
-        assert get_model_url(args, "ref", "/v1/chat/completions") == "http://10.0.0.1:3001/v1/chat/completions"
-
-    def test_get_model_url_fallback(self):
-        """get_model_url should fall back to default router if model not found."""
-        from argparse import Namespace
-
-        from slime.rollout.sglang_rollout import get_model_url
-
-        args = Namespace(
-            sglang_router_ip="10.0.0.1",
-            sglang_router_port=3000,
-            sglang_model_routers={"actor": ("10.0.0.1", 3000)},
-        )
-        assert get_model_url(args, "unknown") == "http://10.0.0.1:3000/generate"
-
-    def test_get_model_url_no_routers(self):
-        """get_model_url should work when sglang_model_routers is not set."""
-        from argparse import Namespace
-
-        from slime.rollout.sglang_rollout import get_model_url
-
-        args = Namespace(
-            sglang_router_ip="10.0.0.1",
-            sglang_router_port=3000,
-        )
-        assert get_model_url(args, "anything") == "http://10.0.0.1:3000/generate"
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v"])
diff --git a/tests/utils/test_trace_utils.py b/tests/utils/test_trace_utils.py
deleted file mode 100644
index 01a6b527b8..0000000000
--- a/tests/utils/test_trace_utils.py
+++ /dev/null
@@ -1,83 +0,0 @@
-import importlib.util
-import sys
-from pathlib import Path
-
-import pytest
-import torch
-
-from slime.utils.trace_utils import build_sglang_meta_trace_attrs, trace_span
-from slime.utils.types import Sample
-
-
-def _load_trace_timeline_viewer_module():
-    module_path = Path(__file__).resolve().parents[2] / "tools" / "trace_timeline_viewer.py"
-    module_name = "test_trace_timeline_viewer_module"
-    sys.modules.pop(module_name, None)
-    spec = importlib.util.spec_from_file_location(module_name, module_path)
-    module = importlib.util.module_from_spec(spec)
-    assert spec is not None
-    assert spec.loader is not None
-    sys.modules[module_name] = module
-    spec.loader.exec_module(module)
-    return module
-
-
-@pytest.mark.unit
-def test_build_sglang_meta_trace_attrs_keeps_standard_and_pd_fields():
-    meta = {
-        "prompt_tokens": 12,
-        "completion_tokens": 7,
-        "cached_tokens": 3,
-        "pd_prefill_forward_duration": 0.125,
-        "pd_decode_transfer_duration": None,
-        "finish_reason": {"type": "stop"},
-        "unused_field": "ignored",
-    }
-
-    assert build_sglang_meta_trace_attrs(meta) == {
-        "prompt_tokens": 12,
-        "completion_tokens": 7,
-        "cached_tokens": 3,
-        "pd_prefill_forward_duration": 0.125,
-        "finish_reason": "stop",
-    }
-
-
-@pytest.mark.unit
-def test_trace_timeline_viewer_omits_virtual_pd_lanes_without_pd_attrs(tmp_path: Path):
-    viewer = _load_trace_timeline_viewer_module()
-    sample = Sample(index=0, prompt="hello")
-
-    with trace_span(sample, "sglang_generate", attrs={"max_new_tokens": 8}) as span:
-        span.update(
-            build_sglang_meta_trace_attrs(
-                {
-                    "prompt_tokens": 4,
-                    "completion_tokens": 2,
-                    "cached_tokens": 1,
-                    "finish_reason": {"type": "stop"},
-                }
-            )
-        )
-
-    pt_path = tmp_path / "rollout.pt"
-    torch.save({"samples": [sample]}, pt_path)
-
-    cache = viewer._build_cache_data(pt_path)
-
-    assert cache["sample_count"] == 1
-    row = cache["rows"][0]
-    assert row["lane_count"] == 1
-    assert row["item_count"] == 1
-    assert row["closed_span_count"] == 1
-
-    item = row["items"][0]
-    assert item["name"] == "sglang_generate"
-    assert item["attrs"]["end_attrs"] == {
-        "prompt_tokens": 4,
-        "completion_tokens": 2,
-        "cached_tokens": 1,
-        "finish_reason": "stop",
-    }
-    assert "[P]" not in item["name"]
-    assert "[D]" not in item["name"]
diff --git a/tools/analyze_profile.py b/tools/analyze_profile.py
deleted file mode 100644
index de511744f2..0000000000
--- a/tools/analyze_profile.py
+++ /dev/null
@@ -1,730 +0,0 @@
-#!/usr/bin/env python3
-"""
-SGLang Decode Profile Analyzer
-==============================
-Analyzes PyTorch profiler traces (.trace.json.gz) from SGLang decode workers.
-
-Usage:
-    python tools/analyze_profile.py --profile-dir profiles/20260303_052303_my_run
-    python tools/analyze_profile.py --profile-dir profiles/20260303_052303_my_run --rank 0
-    python tools/analyze_profile.py --profile-dir profiles/20260303_052303_my_run --all-ranks
-"""
-
-import argparse
-import glob
-import gzip
-import json
-import os
-import sys
-from collections import defaultdict
-from dataclasses import dataclass, field
-
-
-# ─── Color helpers ──────────────────────────────────────────────────────────
-class C:
-    HEADER = "\033[95m"
-    BLUE = "\033[94m"
-    CYAN = "\033[96m"
-    GREEN = "\033[92m"
-    YELLOW = "\033[93m"
-    RED = "\033[91m"
-    BOLD = "\033[1m"
-    UNDERLINE = "\033[4m"
-    END = "\033[0m"
-
-
-def header(text):
-    print(f"\n{C.BOLD}{C.HEADER}{'═' * 80}{C.END}")
-    print(f"{C.BOLD}{C.HEADER}  {text}{C.END}")
-    print(f"{C.BOLD}{C.HEADER}{'═' * 80}{C.END}")
-
-
-def section(text):
-    print(f"\n{C.BOLD}{C.CYAN}── {text} ──{C.END}")
-
-
-def warn(text):
-    print(f"  {C.YELLOW}⚠  {text}{C.END}")
-
-
-def good(text):
-    print(f"  {C.GREEN}✓  {text}{C.END}")
-
-
-def bad(text):
-    print(f"  {C.RED}✗  {text}{C.END}")
-
-
-def bar(pct, width=40, label=""):
-    filled = int(pct / 100 * width)
-    bar_str = "█" * filled + "░" * (width - filled)
-    color = C.GREEN if pct >= 80 else C.YELLOW if pct >= 50 else C.RED
-    return f"  {color}{bar_str}{C.END} {pct:5.1f}%  {label}"
-
-
-# ─── Data structures ───────────────────────────────────────────────────────
-@dataclass
-class KernelInfo:
-    name: str
-    category: str
-    count: int = 0
-    total_dur: float = 0.0
-
-    @property
-    def avg_dur(self):
-        return self.total_dur / max(self.count, 1)
-
-
-@dataclass
-class CudaGraphLaunch:
-    index: int
-    ts: float
-    dur: float
-    tid: int
-
-
-@dataclass
-class TraceAnalysis:
-    rank_name: str = ""
-    # Device
-    gpu_name: str = ""
-    num_gpus: int = 0
-    total_vram_gb: float = 0.0
-    num_sms: int = 0
-    # CUDA/NCCL
-    cuda_version: str = ""
-    nccl_version: str = ""
-    nccl_backend: str = ""
-    world_size: int = 0
-    pg_count: int = 0
-    # Timeline
-    trace_wall_time_us: float = 0.0
-    gpu_active_span_us: float = 0.0
-    gpu_busy_time_us: float = 0.0
-    gpu_util_pct: float = 0.0
-    # Events
-    total_events: int = 0
-    total_kernel_events: int = 0
-    total_kernel_time_us: float = 0.0
-    # Categories
-    kernel_categories: dict[str, KernelInfo] = field(default_factory=dict)
-    top_kernels: list[KernelInfo] = field(default_factory=list)
-    # CUDA Graph
-    cuda_graph_launches: list[CudaGraphLaunch] = field(default_factory=list)
-    cuda_graph_total_cpu_us: float = 0.0
-    # Gaps
-    top_gaps: list[tuple[float, float, float, str]] = field(default_factory=list)  # (dur, start, end, cause)
-    # Decode steps
-    decode_steps: list[dict] = field(default_factory=list)
-    # Per-stream
-    stream_info: dict[int, dict] = field(default_factory=dict)
-    # Communication
-    nccl_kernel_dur_us: float = 0.0
-    deep_ep_dur_us: float = 0.0
-    gloo_total_us: float = 0.0
-    # CPU overhead
-    aten_copy_total_us: float = 0.0
-    aten_copy_count: int = 0
-
-
-def classify_kernel(name: str) -> str:
-    nl = name.lower()
-    if "nccl" in nl:
-        return "NCCL Communication"
-    if "deep_ep" in nl:
-        if "dispatch" in nl:
-            return "DeepEP Dispatch (MoE)"
-        if "combine" in nl:
-            return "DeepEP Combine (MoE)"
-        if "clean" in nl:
-            return "DeepEP Buffer Clean"
-        return "DeepEP (other)"
-    if "flash" in nl and "attn" in nl:
-        return "Flash Attention"
-    if "sparse_attn" in nl:
-        return "Sparse Attention (MLA)"
-    if "paged_mqa" in nl:
-        return "Paged MQA/MLA Logits"
-    if "attention" in nl or "fmha" in nl:
-        return "Attention (other)"
-    if "deep_gemm" in nl or "sm90_fp8_gemm" in nl:
-        return "DeepGEMM (FP8)"
-    if "nvjet" in nl:
-        return "NvJet GEMM"
-    if "gemm" in nl or "cutlass" in nl or "matmul" in nl or "cublas" in nl:
-        return "GEMM/MatMul (other)"
-    if "topk" in nl:
-        return "TopK / MoE Routing"
-    if "quant" in nl:
-        return "Quantization (FP8)"
-    if "rmsnorm" in nl or "layernorm" in nl or "norm" in nl:
-        return "Normalization"
-    if any(k in nl for k in ["triton", "fused"]):
-        return "Fused/Triton Kernels"
-    if "elementwise" in nl or "vectorized_elementwise" in nl:
-        return "Elementwise Ops"
-    if "reduce" in nl or "softmax" in nl:
-        return "Reduce/Softmax"
-    if "embedding" in nl:
-        return "Embedding"
-    if "memcpy" in nl or "memset" in nl:
-        return "Memory Ops"
-    if "index" in nl or "scatter" in nl or "gather" in nl:
-        return "Index/Gather/Scatter"
-    if "cat" in nl and "batched" in nl:
-        return "Concat/Cat"
-    return "Other Compute"
-
-
-def load_trace(filepath: str) -> dict:
-    with gzip.open(filepath, "rt") as f:
-        return json.load(f)
-
-
-def analyze_trace(data: dict, rank_name: str = "") -> TraceAnalysis:
-    result = TraceAnalysis(rank_name=rank_name)
-    events = data["traceEvents"]
-    result.total_events = len(events)
-
-    # ─── Device Properties ──────────────────────────────────────────────
-    dev_props = data.get("deviceProperties", [])
-    if dev_props:
-        result.gpu_name = dev_props[0].get("name", "Unknown")
-        result.num_gpus = len(dev_props)
-        result.total_vram_gb = dev_props[0].get("totalGlobalMem", 0) / (1024**3)
-        result.num_sms = dev_props[0].get("numSms", 0)
-
-    # ─── CUDA/NCCL Info ─────────────────────────────────────────────────
-    cuda_rt = data.get("cuda_runtime_version", 0)
-    result.cuda_version = f"{cuda_rt // 1000}.{(cuda_rt % 1000) // 10}"
-
-    dist_info = data.get("distributedInfo", {})
-    result.nccl_version = dist_info.get("nccl_version", "?")
-    result.nccl_backend = dist_info.get("backend", "?")
-    result.world_size = dist_info.get("world_size", 0)
-    result.pg_count = dist_info.get("pg_count", 0)
-
-    # ─── Kernel Analysis ────────────────────────────────────────────────
-    kernel_events = [e for e in events if e.get("cat") == "kernel" and e.get("ph") == "X"]
-    result.total_kernel_events = len(kernel_events)
-
-    kernel_cats = defaultdict(lambda: {"count": 0, "total_dur": 0.0})
-    kernel_indiv = defaultdict(lambda: {"count": 0, "total_dur": 0.0})
-
-    for e in kernel_events:
-        name = e.get("name", "?")
-        dur = e.get("dur", 0)
-        cat = classify_kernel(name)
-        kernel_cats[cat]["count"] += 1
-        kernel_cats[cat]["total_dur"] += dur
-        kernel_indiv[name]["count"] += 1
-        kernel_indiv[name]["total_dur"] += dur
-
-    result.total_kernel_time_us = sum(v["total_dur"] for v in kernel_cats.values())
-
-    for cat, info in kernel_cats.items():
-        result.kernel_categories[cat] = KernelInfo(
-            name=cat, category=cat, count=info["count"], total_dur=info["total_dur"]
-        )
-
-    for name, info in sorted(kernel_indiv.items(), key=lambda x: -x[1]["total_dur"])[:30]:
-        result.top_kernels.append(
-            KernelInfo(name=name, category=classify_kernel(name), count=info["count"], total_dur=info["total_dur"])
-        )
-
-    # ─── GPU Utilization ────────────────────────────────────────────────
-    gpu_events = [e for e in events if e.get("cat") in ("kernel", "gpu_memcpy", "gpu_memset") and e.get("ph") == "X"]
-    intervals = sorted([(e["ts"], e["ts"] + e["dur"]) for e in gpu_events])
-
-    if intervals:
-        merged = []
-        for s, e_end in intervals:
-            if merged and s <= merged[-1][1]:
-                merged[-1] = (merged[-1][0], max(merged[-1][1], e_end))
-            else:
-                merged.append((s, e_end))
-
-        result.gpu_busy_time_us = sum(e - s for s, e in merged)
-        result.gpu_active_span_us = merged[-1][1] - merged[0][0]
-        result.gpu_util_pct = result.gpu_busy_time_us / result.gpu_active_span_us * 100
-
-        # Trace wall time
-        all_x = [(e["ts"], e["ts"] + e["dur"]) for e in events if e.get("ph") == "X" and "ts" in e and "dur" in e]
-        if all_x:
-            result.trace_wall_time_us = max(t[1] for t in all_x) - min(t[0] for t in all_x)
-
-        # GPU idle gaps
-        gaps = []
-        for i in range(1, len(merged)):
-            gap = merged[i][0] - merged[i - 1][1]
-            if gap > 0:
-                gaps.append((gap, merged[i - 1][1], merged[i][0]))
-        gaps.sort(reverse=True)
-
-        # Identify cause of top gaps
-        cpu_ops = [
-            e for e in events if e.get("cat") in ("cpu_op", "user_annotation", "cuda_runtime") and e.get("ph") == "X"
-        ]
-        for gap, gap_start, gap_end in gaps[:15]:
-            cause = "unknown"
-            for e in cpu_ops:
-                e_start = e.get("ts", 0)
-                e_end = e_start + e.get("dur", 0)
-                if e_start < gap_end and e_end > gap_start:
-                    cause = e.get("name", "?")
-                    break
-            result.top_gaps.append((gap, gap_start, gap_end, cause))
-
-    # ─── Per-stream analysis ────────────────────────────────────────────
-    stream_events = defaultdict(list)
-    for e in gpu_events:
-        stream_events[e.get("tid")].append(e)
-
-    for tid, evts in stream_events.items():
-        total_dur = sum(e.get("dur", 0) for e in evts)
-        min_ts = min(e["ts"] for e in evts)
-        max_ts = max(e["ts"] + e.get("dur", 0) for e in evts)
-        span = max_ts - min_ts
-        result.stream_info[tid] = {
-            "count": len(evts),
-            "total_dur": total_dur,
-            "span": span,
-            "util": total_dur / span * 100 if span > 0 else 0,
-        }
-
-    # ─── CUDA Graph analysis ───────────────────────────────────────────
-    cg_launches = sorted(
-        [e for e in events if e.get("name") == "cudaGraphLaunch" and e.get("cat") == "cuda_runtime"],
-        key=lambda x: x["ts"],
-    )
-    for i, e in enumerate(cg_launches):
-        result.cuda_graph_launches.append(CudaGraphLaunch(index=i, ts=e["ts"], dur=e["dur"], tid=e.get("tid", 0)))
-    result.cuda_graph_total_cpu_us = sum(e["dur"] for e in cg_launches)
-
-    # Group into decode steps (3 launches per step)
-    for i in range(0, len(cg_launches), 3):
-        group = cg_launches[i : i + 3]
-        if len(group) == 3:
-            step_start = group[0]["ts"]
-            step_end = group[2]["ts"] + group[2]["dur"]
-            result.decode_steps.append(
-                {
-                    "step": i // 3,
-                    "span_us": step_end - step_start,
-                    "launches": [g["dur"] for g in group],
-                    "ts_start": step_start,
-                    "ts_end": step_end,
-                }
-            )
-
-    # ─── Communication ──────────────────────────────────────────────────
-    result.nccl_kernel_dur_us = sum(
-        e.get("dur", 0) for e in events if e.get("cat") == "kernel" and "nccl" in e.get("name", "").lower()
-    )
-    result.deep_ep_dur_us = sum(
-        e.get("dur", 0) for e in events if e.get("cat") == "kernel" and "deep_ep" in e.get("name", "")
-    )
-    result.gloo_total_us = sum(e.get("dur", 0) for e in events if "gloo" in e.get("name", "").lower())
-
-    # ─── CPU overhead ───────────────────────────────────────────────────
-    copy_events = [e for e in events if e.get("cat") == "cpu_op" and e.get("name") == "aten::copy_"]
-    result.aten_copy_total_us = sum(e.get("dur", 0) for e in copy_events)
-    result.aten_copy_count = len(copy_events)
-
-    return result
-
-
-# ─── Pretty print ──────────────────────────────────────────────────────────
-def print_analysis(r: TraceAnalysis):
-    header(f"Profile Analysis: {r.rank_name}")
-
-    # ── Config ──────────────────────────────────────────────────────────
-    section("Hardware & Config")
-    print(f"  GPU:            {C.BOLD}{r.gpu_name}{C.END} × {r.num_gpus} (visible)")
-    print(f"  VRAM per GPU:   {r.total_vram_gb:.0f} GB")
-    print(f"  SMs per GPU:    {r.num_sms}")
-    print(f"  CUDA Version:   {r.cuda_version}")
-    print(f"  NCCL:           {r.nccl_version} (backend: {r.nccl_backend})")
-    print(f"  World Size:     {r.world_size}")
-    print(f"  Process Groups: {r.pg_count}")
-    print(f"  Total Events:   {r.total_events:,}")
-
-    # ── Timeline ────────────────────────────────────────────────────────
-    section("Timeline")
-    print(f"  Trace Wall Time:  {r.trace_wall_time_us / 1000:.1f} ms")
-    print(f"  GPU Active Span:  {r.gpu_active_span_us / 1000:.1f} ms")
-    print(f"  GPU Busy Time:    {r.gpu_busy_time_us / 1000:.1f} ms")
-    print(f"  GPU Idle (bubbles): {(r.gpu_active_span_us - r.gpu_busy_time_us) / 1000:.1f} ms")
-    print()
-    print(f"  {C.BOLD}GPU Utilization:{C.END}")
-    print(bar(r.gpu_util_pct, label="GPU busy"))
-    print(bar(100 - r.gpu_util_pct, label="GPU idle (bubbles)"))
-
-    # ── Kernel Breakdown ────────────────────────────────────────────────
-    section("GPU Kernel Time Breakdown")
-    total = r.total_kernel_time_us
-    print(f"  Total kernel time: {total / 1000:.1f} ms ({r.total_kernel_events:,} kernel launches)\n")
-
-    sorted_cats = sorted(r.kernel_categories.values(), key=lambda x: -x.total_dur)
-    for ki in sorted_cats:
-        pct = ki.total_dur / total * 100
-        filled = int(pct / 100 * 30)
-        bar_str = "█" * filled + "░" * (30 - filled)
-        color = C.CYAN
-        if "NCCL" in ki.name or "DeepEP" in ki.name:
-            color = C.YELLOW
-        elif "GEMM" in ki.name or "NvJet" in ki.name or "DeepGEMM" in ki.name:
-            color = C.GREEN
-        elif "Attention" in ki.name or "MLA" in ki.name or "MQA" in ki.name:
-            color = C.BLUE
-        print(f"  {color}{bar_str}{C.END} {pct:5.1f}%  {ki.total_dur / 1000:8.1f} ms  n={ki.count:5d}  {ki.name}")
-
-    # ── Top Kernels ─────────────────────────────────────────────────────
-    section("Top 20 Individual Kernels (by total GPU time)")
-    print(f"  {'%':>5s}  {'Total(ms)':>9s}  {'Avg(us)':>8s}  {'Count':>6s}  Kernel")
-    print(f"  {'─' * 5}  {'─' * 9}  {'─' * 8}  {'─' * 6}  {'─' * 50}")
-    for ki in r.top_kernels[:20]:
-        pct = ki.total_dur / total * 100
-        print(f"  {pct:5.1f}  {ki.total_dur / 1000:9.2f}  {ki.avg_dur:8.1f}  {ki.count:6d}  {ki.name[:90]}")
-
-    # ── CUDA Graph ──────────────────────────────────────────────────────
-    section("CUDA Graph Analysis")
-    print(f"  Total cudaGraphLaunch calls: {C.BOLD}{len(r.cuda_graph_launches)}{C.END}")
-    print(f"  Total CPU time in cudaGraphLaunch: {C.BOLD}{r.cuda_graph_total_cpu_us / 1000:.1f} ms{C.END}")
-    print()
-
-    if r.cuda_graph_launches:
-        # Categorize launches by duration
-        small = [launch for launch in r.cuda_graph_launches if launch.dur < 1000]
-        medium = [launch for launch in r.cuda_graph_launches if 1000 <= launch.dur < 5000]
-        large = [launch for launch in r.cuda_graph_launches if launch.dur >= 5000]
-
-        print("  Launch size distribution:")
-        print(
-            f"    Small  (<1ms):  {len(small):3d} launches, total {sum(launch.dur for launch in small) / 1000:.1f} ms"
-        )
-        print(
-            f"    Medium (1-5ms): {len(medium):3d} launches, total {sum(launch.dur for launch in medium) / 1000:.1f} ms"
-        )
-        print(
-            f"    {C.RED}Large  (>5ms):  {len(large):3d} launches, total {sum(launch.dur for launch in large) / 1000:.1f} ms{C.END}"
-        )
-
-    if r.decode_steps:
-        print("\n  Decode Steps (3 cudaGraphLaunch per step):")
-        print(f"  {'Step':>4s}  {'Span(ms)':>8s}  {'L1(us)':>7s}  {'L2(us)':>7s}  {'L3(us)':>7s}  Note")
-        print(f"  {'─' * 4}  {'─' * 8}  {'─' * 7}  {'─' * 7}  {'─' * 7}  {'─' * 30}")
-        for step in r.decode_steps:
-            launches = step["launches"]
-            note = ""
-            if any(val > 10000 for val in launches):
-                note = f"{C.RED}← large launch stall{C.END}"
-            print(
-                f"  {step['step']:4d}  {step['span_us'] / 1000:8.1f}  "
-                f"{launches[0]:7.0f}  {launches[1]:7.0f}  {launches[2]:7.0f}  {note}"
-            )
-
-        avg_span = sum(s["span_us"] for s in r.decode_steps) / len(r.decode_steps)
-        print(f"\n  Average decode step span: {C.BOLD}{avg_span / 1000:.1f} ms{C.END}")
-        print(f"  Estimated decode throughput: {C.BOLD}{1e6 / avg_span:.0f} tokens/s per GPU{C.END}")
-
-    # ── CUDA Graph Location in Timeline ─────────────────────────────────
-    section("CUDA Graph Locations in Timeline")
-    if r.cuda_graph_launches and r.gpu_active_span_us > 0:
-        base_ts = min(launch.ts for launch in r.cuda_graph_launches)
-        span = r.gpu_active_span_us
-        print(f"  Timeline (0 = first GPU activity, total span = {span / 1000:.1f} ms):")
-        print()
-        # ASCII timeline
-        WIDTH = 70
-        timeline = [" "] * WIDTH
-        for launch in r.cuda_graph_launches:
-            pos = int((launch.ts - base_ts) / span * WIDTH)
-            pos = min(pos, WIDTH - 1)
-            if launch.dur >= 5000:
-                timeline[pos] = f"{C.RED}▓{C.END}"
-            elif launch.dur >= 1000:
-                timeline[pos] = f"{C.YELLOW}▒{C.END}"
-            else:
-                timeline[pos] = f"{C.GREEN}░{C.END}"
-        print(f"  |{''.join(timeline)}|")
-        print(f"  0ms{' ' * (WIDTH - 8)}{span / 1000:.0f}ms")
-        print(
-            f"  Legend: {C.GREEN}░{C.END}=small(<1ms)  {C.YELLOW}▒{C.END}=medium(1-5ms)  {C.RED}▓{C.END}=large(>5ms)"
-        )
-
-    # ── GPU Idle Gaps ───────────────────────────────────────────────────
-    section("Top GPU Idle Gaps (Bubbles)")
-    if r.top_gaps:
-        print(f"  {'#':>3s}  {'Gap(us)':>8s}  {'Cause'}")
-        print(f"  {'─' * 3}  {'─' * 8}  {'─' * 50}")
-        for i, (gap, _start, _end, cause) in enumerate(r.top_gaps[:10]):
-            color = C.RED if gap > 3000 else C.YELLOW if gap > 1000 else ""
-            end_c = C.END if color else ""
-            print(f"  {i + 1:3d}  {color}{gap:8.0f}{end_c}  {cause}")
-
-    # ── Per-Stream ──────────────────────────────────────────────────────
-    section("Per-Stream GPU Activity")
-    print(f"  {'Stream':>8s}  {'Events':>7s}  {'Time(ms)':>9s}  {'Span(ms)':>9s}  {'Util%':>6s}  Role")
-    print(f"  {'─' * 8}  {'─' * 7}  {'─' * 9}  {'─' * 9}  {'─' * 6}  {'─' * 20}")
-    for tid in sorted(r.stream_info.keys()):
-        info = r.stream_info[tid]
-        role = ""
-        if info["count"] > 50000:
-            role = "← main compute"
-        elif info["count"] > 5000:
-            role = "← secondary compute"
-        elif info["util"] < 2:
-            role = "← auxiliary"
-        print(
-            f"  {tid:8d}  {info['count']:7d}  {info['total_dur'] / 1000:9.1f}  "
-            f"{info['span'] / 1000:9.1f}  {info['util']:6.1f}  {role}"
-        )
-
-    # ── Communication ───────────────────────────────────────────────────
-    section("Communication Overhead")
-    comm_total = r.nccl_kernel_dur_us + r.deep_ep_dur_us
-    print(f"  NCCL AllGather (GPU kernel): {r.nccl_kernel_dur_us / 1000:.1f} ms")
-    print(f"  DeepEP (MoE all-to-all):     {r.deep_ep_dur_us / 1000:.1f} ms")
-    print(f"  Gloo Broadcast (CPU):        {r.gloo_total_us / 1000:.1f} ms")
-    print(f"  Communication / Total Kernel: {comm_total / total * 100:.1f}%")
-
-    # ── Performance Bottleneck Summary ──────────────────────────────────
-    header("🔍 Performance Bottleneck Analysis")
-
-    issues = []
-
-    # 1. GPU utilization
-    if r.gpu_util_pct < 90:
-        issues.append(
-            (
-                "GPU Utilization",
-                f"GPU utilization is {r.gpu_util_pct:.1f}% — {100 - r.gpu_util_pct:.1f}% idle bubbles "
-                f"({(r.gpu_active_span_us - r.gpu_busy_time_us) / 1000:.1f} ms wasted)",
-                "high" if r.gpu_util_pct < 80 else "medium",
-            )
-        )
-
-    # 2. CUDA Graph launch overhead
-    large_launches = [launch for launch in r.cuda_graph_launches if launch.dur >= 5000]
-    if large_launches:
-        avg_large = sum(launch.dur for launch in large_launches) / len(large_launches)
-        issues.append(
-            (
-                "CUDA Graph Launch Stalls",
-                f"{len(large_launches)} cudaGraphLaunch calls take >{5}ms (avg {avg_large / 1000:.1f}ms). "
-                f"Total: {sum(launch.dur for launch in large_launches) / 1000:.1f}ms CPU blocked. "
-                f"This is the {C.BOLD}#1 source of GPU bubbles{C.END}.",
-                "high",
-            )
-        )
-
-    # 3. DeepEP dispatch latency
-    dep_cat = r.kernel_categories.get("DeepEP Dispatch (MoE)")
-    if dep_cat and dep_cat.total_dur / total > 0.15:
-        issues.append(
-            (
-                "DeepEP Dispatch Dominance",
-                f"DeepEP dispatch takes {dep_cat.total_dur / 1000:.1f}ms ({dep_cat.total_dur / total * 100:.1f}% of kernel time). "
-                f"MoE expert parallelism communication is a major cost.",
-                "high",
-            )
-        )
-
-    # 4. MoE routing overhead
-    topk_cat = r.kernel_categories.get("TopK / MoE Routing")
-    if topk_cat and topk_cat.total_dur / total > 0.05:
-        issues.append(
-            (
-                "MoE Routing Overhead",
-                f"TopK routing takes {topk_cat.total_dur / 1000:.1f}ms ({topk_cat.total_dur / total * 100:.1f}%). "
-                f"topk_transform_decode_kernel is expensive at {topk_cat.avg_dur:.0f}us avg.",
-                "medium",
-            )
-        )
-
-    # 5. CPU copy overhead
-    if r.aten_copy_total_us > 100000:
-        issues.append(
-            (
-                "CPU aten::copy_ Overhead",
-                f"aten::copy_ takes {r.aten_copy_total_us / 1000:.1f}ms CPU time ({r.aten_copy_count} calls). "
-                f"Likely dtype conversions (BF16→FP8) stalling the CPU.",
-                "medium",
-            )
-        )
-
-    # 6. Small kernel launches
-    small_kernels = sum(1 for k in r.top_kernels if k.avg_dur < 3)
-    if small_kernels > 5:
-        issues.append(
-            (
-                "Many Tiny Kernels",
-                f"{small_kernels} of top-30 kernels have avg duration <3us. "
-                f"Launch overhead may exceed compute. Consider kernel fusion.",
-                "low",
-            )
-        )
-
-    for issue_name, desc, severity in issues:
-        color = C.RED if severity == "high" else C.YELLOW if severity == "medium" else C.CYAN
-        sev_label = {"high": "HIGH", "medium": "MED", "low": "LOW"}[severity]
-        print(f"\n  {color}[{sev_label}]{C.END} {C.BOLD}{issue_name}{C.END}")
-        print(f"       {desc}")
-
-    # ── Optimization Recommendations ────────────────────────────────────
-    header("💡 Optimization Recommendations")
-
-    recs = []
-
-    # Based on detected issues
-    if any("CUDA Graph Launch" in i[0] for i in issues):
-        recs.append(
-            (
-                "Reduce CUDA Graph Re-capture / Launch Latency",
-                [
-                    "The large cudaGraphLaunch (13-16ms) occurs once per decode step — this is the MoE expert "
-                    "portion that runs OUTSIDE the CUDA graph (DeepEP dispatch/combine + NCCL allgather).",
-                    "The 3-launch pattern per step = (1) pre-MoE graph, (2) post-MoE graph, (3) MoE-expert graph.",
-                    "Optimization: try increasing decode batch size to amortize graph launch overhead per token.",
-                    "Check if `--sglang-disable-cuda-graph` helps isolate whether the overhead is in graph "
-                    "management vs. actual compute.",
-                    "Consider padding batch sizes to avoid frequent graph re-capture for different sizes.",
-                ],
-            )
-        )
-
-    if any("DeepEP" in i[0] for i in issues):
-        recs.append(
-            (
-                "Optimize MoE Expert Parallelism (DeepEP)",
-                [
-                    f"DeepEP dispatch+combine = {r.deep_ep_dur_us / 1000:.1f}ms/step = "
-                    f"{r.deep_ep_dur_us / total * 100:.1f}% of GPU time — this is all-to-all expert communication.",
-                    "dispatch (76.8us avg × 1556 calls) dominates over combine (29.4us avg).",
-                    "Consider: reduce number of MoE layers, or reduce EP degree if not fully utilizing all experts.",
-                    "Ensure NVLink/NVSwitch bandwidth is saturated (H100 should be 900 GB/s bidirectional).",
-                    "Check if low-latency mode for DeepEP is enabled — the clean_buffer kernel (700us avg) suggests "
-                    "low-latency mode is active.",
-                ],
-            )
-        )
-
-    if any("TopK" in i[0] for i in issues):
-        recs.append(
-            (
-                "Optimize MoE TopK Routing",
-                [
-                    "topk_transform_decode_kernel at 49us avg is expensive for decode (small batch).",
-                    "Consider using a more efficient routing algorithm or fusing TopK with dispatch.",
-                ],
-            )
-        )
-
-    recs.append(
-        (
-            "General Decode Throughput",
-            [
-                (
-                    f"Current: ~{1e6 / (sum(s['span_us'] for s in r.decode_steps) / max(len(r.decode_steps), 1)):.0f} tokens/s/GPU "
-                    f"(avg {sum(s['span_us'] for s in r.decode_steps) / max(len(r.decode_steps), 1) / 1000:.1f}ms/step)."
-                    if r.decode_steps
-                    else "No decode step data."
-                ),
-                "Increase batch size to improve GPU SM occupancy — many kernels are memory-bound at small batch.",
-                "Speculative decoding could help if generation is latency-bound.",
-                "Verify `--sglang-mem-fraction-static` is set high enough for large KV cache.",
-            ],
-        )
-    )
-
-    for i, (title, items) in enumerate(recs):
-        print(f"\n  {C.BOLD}{i + 1}. {title}{C.END}")
-        for item in items:
-            print(f"     • {item}")
-
-
-def print_cross_rank_summary(analyses: list[TraceAnalysis]):
-    header("Cross-Rank Comparison")
-    print(
-        f"  {'Rank':<25s}  {'Span(ms)':>9s}  {'Busy(ms)':>9s}  {'Util%':>6s}  "
-        f"{'Kernel(ms)':>10s}  {'DeepEP(ms)':>10s}  {'GEMM(ms)':>9s}  {'NCCL(ms)':>9s}"
-    )
-    print(f"  {'─' * 25}  {'─' * 9}  {'─' * 9}  {'─' * 6}  {'─' * 10}  {'─' * 10}  {'─' * 9}  {'─' * 9}")
-
-    for r in analyses:
-        gemm_dur = sum(
-            ki.total_dur
-            for ki in r.kernel_categories.values()
-            if "GEMM" in ki.name or "NvJet" in ki.name or "DeepGEMM" in ki.name
-        )
-        print(
-            f"  {r.rank_name:<25s}  {r.gpu_active_span_us / 1000:9.1f}  {r.gpu_busy_time_us / 1000:9.1f}  "
-            f"{r.gpu_util_pct:6.1f}  {r.total_kernel_time_us / 1000:10.1f}  "
-            f"{r.deep_ep_dur_us / 1000:10.1f}  {gemm_dur / 1000:9.1f}  {r.nccl_kernel_dur_us / 1000:9.1f}"
-        )
-
-    utils = [r.gpu_util_pct for r in analyses]
-    spans = [r.gpu_active_span_us for r in analyses]
-    print(f"\n  Utilization range: {min(utils):.1f}% ~ {max(utils):.1f}% (spread: {max(utils) - min(utils):.1f}%)")
-    print(f"  Span range: {min(spans) / 1000:.1f}ms ~ {max(spans) / 1000:.1f}ms")
-
-    if max(utils) - min(utils) > 5:
-        warn("Significant load imbalance detected across ranks (>5% spread).")
-    else:
-        good("Load balance across ranks is reasonable (<5% spread).")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="Analyze SGLang decode profile traces")
-    parser.add_argument("--profile-dir", type=str, required=True, help="Directory containing .trace.json.gz files")
-    parser.add_argument("--rank", type=int, default=None, help="Specific rank to analyze (default: first file)")
-    parser.add_argument("--all-ranks", action="store_true", help="Analyze all ranks and show comparison")
-    parser.add_argument("--top-n", type=int, default=20, help="Number of top kernels to show")
-    args = parser.parse_args()
-
-    files = sorted(glob.glob(os.path.join(args.profile_dir, "*.trace.json.gz")))
-    if not files:
-        print(f"No .trace.json.gz files found in {args.profile_dir}")
-        sys.exit(1)
-
-    print(f"Found {len(files)} trace files in {args.profile_dir}")
-
-    if args.all_ranks:
-        analyses = []
-        for fpath in files:
-            basename = os.path.basename(fpath)
-            # Extract rank info
-            parts = basename.split("-")
-            rank_name = "-".join(parts[1:]).replace(".trace.json.gz", "")
-            print(f"  Loading {rank_name}...", end=" ", flush=True)
-            data = load_trace(fpath)
-            analysis = analyze_trace(data, rank_name)
-            analyses.append(analysis)
-            print("done")
-        # Full analysis of first rank
-        print_analysis(analyses[0])
-        # Cross-rank comparison
-        print_cross_rank_summary(analyses)
-    else:
-        if args.rank is not None:
-            target = f"TP-{args.rank}"
-            matching = [f for f in files if target in f]
-            if not matching:
-                print(f"No file found for rank {args.rank}")
-                sys.exit(1)
-            fpath = matching[0]
-        else:
-            fpath = files[0]
-
-        basename = os.path.basename(fpath)
-        rank_name = "-".join(basename.split("-")[1:]).replace(".trace.json.gz", "")
-        print(f"Analyzing: {basename}")
-
-        data = load_trace(fpath)
-        analysis = analyze_trace(data, rank_name)
-        print_analysis(analysis)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/tools/convert_fsdp_to_hf.py b/tools/convert_fsdp_to_hf.py
new file mode 100644
index 0000000000..ff1ee30ea1
--- /dev/null
+++ b/tools/convert_fsdp_to_hf.py
@@ -0,0 +1,185 @@
+import argparse
+import os
+import pickle
+import shutil
+import time
+
+import torch
+import torch.distributed.checkpoint as dist_cp
+from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForImageTextToText
+from typing_extensions import override
+
+
+class UnpicklerWrapper(pickle.Unpickler):
+    @override
+    def find_class(self, mod_name, name):
+        class DummyClass:
+            def __init__(self, *args, **kwargs):
+                pass
+
+        if mod_name.startswith("megatron") or mod_name.startswith("glm"):
+            return DummyClass
+        return super().find_class(mod_name, name)
+
+
+class WrappedStorageReader(dist_cp.FileSystemReader):
+    @override
+    def read_metadata(self):
+        path = self.fs.concat_path(self.path, ".metadata")
+        with self.fs.create_stream(path, "rb") as metadata_file:
+            metadata = UnpicklerWrapper(metadata_file).load()
+        if getattr(metadata, "storage_meta", None) is None:
+            metadata.storage_meta = dist_cp.StorageMeta()
+        metadata.storage_meta.load_id = self.load_id
+        if metadata.planner_data is None:
+            metadata.planner_data = {}
+        return metadata
+
+
+class EmptyStateDictLoadPlanner(dist_cp.default_planner.DefaultLoadPlanner):
+    @override
+    def set_up_planner(
+        self,
+        state_dict: dist_cp.metadata.STATE_DICT_TYPE,
+        metadata: dist_cp.metadata.Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        for k, v in metadata.state_dict_metadata.items():
+            if "optimizer" in k:
+                continue
+            print(f"find {k} in torch_dist ckpt")
+            if isinstance(v, dist_cp.metadata.TensorStorageMetadata):
+                v = torch.empty(v.size, dtype=v.properties.dtype)  # type: ignore[assignment]
+            state_dict[k] = v
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+
+def _detect_model_dir(input_dir: str) -> str:
+    model_dir = os.path.join(input_dir, "model")
+    return model_dir if os.path.isdir(model_dir) else input_dir
+
+
+def _load_fsdp_state_dict(input_dir: str) -> dict[str, torch.Tensor]:
+    state_dict: dict[str, torch.Tensor] = {}
+    dist_cp.state_dict_loader._load_state_dict(
+        state_dict,
+        storage_reader=WrappedStorageReader(input_dir),
+        planner=EmptyStateDictLoadPlanner(),
+        no_dist=True,
+    )
+    return state_dict
+
+
+def _get_candidate_prefixes(keys: list[str]) -> list[str]:
+    predefined = [
+        "model_state.model.",
+        "model_state.",
+        "model.",
+        "module.",
+        "",
+    ]
+
+    detected: set[str] = set()
+    for key in keys:
+        for prefix in predefined:
+            if prefix and key.startswith(prefix):
+                detected.add(prefix)
+
+    # Always keep empty string as a fall back option for exact match.
+    detected.add("")
+    # Preserve predefined order while keeping only detected prefixes.
+    return [p for p in predefined if p in detected]
+
+
+def _strip_best_prefix(keys: list[str], target_keys: set[str]) -> tuple[str, int]:
+    best_prefix = ""
+    best_match = -1
+
+    for prefix in _get_candidate_prefixes(keys):
+        mapped_keys = {k.removeprefix(prefix) for k in keys}
+        match_count = len(mapped_keys & target_keys)
+        if match_count > best_match:
+            best_match = match_count
+            best_prefix = prefix
+
+    return best_prefix, best_match
+
+
+def _build_hf_model(config: AutoConfig) -> torch.nn.Module:
+    print(f"Detected model type: {config.model_type}")
+    model_cls = AutoModelForImageTextToText if hasattr(config, "vision_config") else AutoModelForCausalLM
+    print(f"Loaded with {model_cls.__name__}")
+    return model_cls.from_config(config, trust_remote_code=True)
+
+
+def _convert_fsdp_to_hf(
+    origin_hf_dir: str,
+    input_dir: str,
+    output_dir: str,
+) -> None:
+    print(f"loading FSDP model from {input_dir}")
+    t = time.time()
+    state_dict = _load_fsdp_state_dict(input_dir)
+    print(f"FSDP model loaded in {time.time()-t:.2f} sec.")
+
+    tensor_items = {k: v for k, v in state_dict.items() if isinstance(v, torch.Tensor)}
+
+    config = AutoConfig.from_pretrained(origin_hf_dir, trust_remote_code=True)
+    hf_model = _build_hf_model(config)
+    target_keys = set(hf_model.state_dict().keys())
+
+    best_prefix, best_match = _strip_best_prefix(list(tensor_items.keys()), target_keys)
+    total_keys = len(tensor_items)
+
+    print(f"Using prefix '{best_prefix}' for key mapping. " f"Matched {best_match}/{total_keys} parameter keys.")
+
+    model_state = {k.removeprefix(best_prefix): v for k, v in tensor_items.items()}
+
+    if not model_state:
+        raise ValueError(
+            "No model weights found in checkpoint. "
+            "Please pass the checkpoint directory (e.g. iter_xxx or iter_xxx/model)."
+        )
+
+    missing, unexpected = hf_model.load_state_dict(model_state, strict=False)
+    print(f"Missing keys: {missing}\nUnexpected keys: {unexpected}")
+
+    os.makedirs(output_dir, exist_ok=True)
+    hf_model.save_pretrained(output_dir, safe_serialization=True)
+    print(f"Model weights saved to {output_dir}")
+
+
+def copy_assets(origin_hf_dir: str, output_dir: str) -> None:
+    for filename in os.listdir(origin_hf_dir):
+        if filename == "model.safetensors.index.json" or filename.endswith(".safetensors"):
+            continue
+        origin_filename = os.path.join(origin_hf_dir, filename)
+        if not os.path.isfile(origin_filename):
+            print(f"Skip {filename}, not a file.")
+            continue
+        src, dst = origin_filename, os.path.join(output_dir, filename)
+        print(f"copy from {src} to {dst}")
+        shutil.copy(src, dst)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input-dir", type=str, required=True)
+    parser.add_argument("--output-dir", type=str, required=True)
+    parser.add_argument(
+        "--origin-hf-dir",
+        type=str,
+        required=True,
+        help="The original Hugging Face model directory to load config/tokenizer assets.",
+    )
+    parser.add_argument(
+        "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists."
+    )
+    args = parser.parse_args()
+
+    if os.path.exists(args.output_dir) and not args.force:
+        raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.")
+
+    model_dir = _detect_model_dir(args.input_dir)
+    _convert_fsdp_to_hf(args.origin_hf_dir, model_dir, args.output_dir)
+    copy_assets(args.origin_hf_dir, args.output_dir)
diff --git a/tools/convert_hf_to_fp8.py b/tools/convert_hf_to_fp8.py
index 5294140d5a..98eebce01b 100644
--- a/tools/convert_hf_to_fp8.py
+++ b/tools/convert_hf_to_fp8.py
@@ -134,11 +134,6 @@ def process_file(input_path, output_path, filename, strategy, block_size, result
             and "lm_head" not in key
             and "eh_proj" not in key
             and "weights_proj" not in key
-            and "conv1d" not in key
-            and "A_log" not in key
-            and "dt_bias" not in key
-            and "in_proj_a" not in key
-            and "in_proj_b" not in key
         ):
             qw, s = quant_fp8(weights[key], strategy, block_size)
             q_weights[key] = qw
diff --git a/tools/convert_torch_dist_to_hf.py b/tools/convert_torch_dist_to_hf.py
index 8049d77437..2ddd1b7b1b 100644
--- a/tools/convert_torch_dist_to_hf.py
+++ b/tools/convert_torch_dist_to_hf.py
@@ -103,20 +103,18 @@ def get_named_params(args, state_dict):
         yield from get_layer_param(args, name, param)
 
 
-def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_size=None, origin_hf_dir=None):
+def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_size=None):
     print(f"start saving to {output_dir}")
     os.makedirs(output_dir, exist_ok=True)
     # 2GB
     current_size = 0
     total_size = 0
     modeltensors = [{}]
-    converted_names = set()
     for name, param in get_named_params(args, state_dict):
         if vocab_size:
             param = remove_padding(name, param, vocab_size)
         converted_named_tensors = convert_to_hf(args, model_name, name, param)
         for converted_name, converted_param in converted_named_tensors:
-            converted_names.add(converted_name)
             tensor_size = converted_param.numel() * converted_param.element_size()
             if tensor_size + current_size > chunk_size:
                 modeltensors.append({})
@@ -125,24 +123,6 @@ def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_siz
             current_size += tensor_size
             total_size += tensor_size
 
-    if origin_hf_dir is not None:
-        safetensors_files = [f for f in os.listdir(origin_hf_dir) if f.endswith(".safetensors")]
-        for filename in safetensors_files:
-            with safetensors.safe_open(os.path.join(origin_hf_dir, filename), framework="pt", device="cpu") as f:
-                for k in f.keys():
-                    if k not in converted_names:
-                        converted_name = k
-                        print(f"add {k} from origin hf checkpoint")
-                        converted_param = f.get_tensor(k)
-                        converted_names.add(k)
-                        tensor_size = converted_param.numel() * converted_param.element_size()
-                        if tensor_size + current_size > chunk_size:
-                            modeltensors.append({})
-                            current_size = 0
-                        modeltensors[-1][converted_name] = converted_param
-                        current_size += tensor_size
-                    total_size += tensor_size
-
     metadata = {"metadata": {"total_size": total_size}, "weight_map": {}}
 
     num_files = len(modeltensors)
@@ -189,9 +169,6 @@ def copy_assets(origin_hf_dir, output_dir):
     parser.add_argument(
         "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists."
     )
-    parser.add_argument(
-        "-a", "--add-missing-from-origin-hf", action="store_true", help="Add missing weights from origin hf checkpoint"
-    )
     parser.add_argument(
         "--chunk-size",
         type=int,
@@ -230,15 +207,7 @@ def copy_assets(origin_hf_dir, output_dir):
     )
     print(f"model loaded in {time.time()-t:.2f} sec.")
 
-    save_tensors(
-        megatron_args,
-        args.model_name,
-        state_dict,
-        args.output_dir,
-        args.chunk_size,
-        args.vocab_size,
-        args.origin_hf_dir if args.add_missing_from_origin_hf else None,
-    )
+    save_tensors(megatron_args, args.model_name, state_dict, args.output_dir, args.chunk_size, args.vocab_size)
 
     if args.origin_hf_dir:
         copy_assets(args.origin_hf_dir, args.output_dir)
diff --git a/tools/convert_torch_dist_to_hf_bridge.py b/tools/convert_torch_dist_to_hf_bridge.py
index 798503f218..3412e73e69 100644
--- a/tools/convert_torch_dist_to_hf_bridge.py
+++ b/tools/convert_torch_dist_to_hf_bridge.py
@@ -4,8 +4,6 @@
 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.
@@ -51,7 +49,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 = patch_auto_bridge_hf_config(AutoBridge.from_hf_pretrained(args.origin_hf_dir, trust_remote_code=True))
+    bridge = 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.
diff --git a/tools/convert_torch_dist_to_hf_parallel.py b/tools/convert_torch_dist_to_hf_parallel.py
deleted file mode 100644
index 763254d42c..0000000000
--- a/tools/convert_torch_dist_to_hf_parallel.py
+++ /dev/null
@@ -1,578 +0,0 @@
-import argparse
-import json
-import multiprocessing
-import os
-import pickle
-import re
-import shutil
-import threading
-import time
-from concurrent.futures import ThreadPoolExecutor, as_completed
-
-import safetensors.torch
-import torch
-import torch.distributed.checkpoint as dist_cp
-from tqdm import tqdm
-from transformers import AutoConfig
-from typing_extensions import override
-
-from slime.backends.megatron_utils.megatron_to_hf import convert_to_hf, remove_padding
-
-
-class DummyClass:
-    def __init__(self, *args, **kwargs):
-        pass
-
-
-class UnpicklerWrapper(pickle.Unpickler):
-    @override
-    def find_class(self, mod_name, name):
-        if mod_name.startswith("megatron") or mod_name.startswith("glm"):
-            return DummyClass
-        return super().find_class(mod_name, name)
-
-
-pickle.Unpickler = UnpicklerWrapper
-
-
-class WrappedStorageReader(dist_cp.FileSystemReader):
-    def __init__(self, path, load_id=None, max_workers=None):
-        super().__init__(path, load_id)
-        self.max_workers = max_workers
-        self._read_cache = {} if max_workers and max_workers > 1 else None
-        self._cache_lock = threading.Lock() if max_workers and max_workers > 1 else None
-
-    @override
-    def read_metadata(self):
-        path = self.fs.concat_path(self.path, ".metadata")
-        with self.fs.create_stream(path, "rb") as metadata_file:
-            metadata = UnpicklerWrapper(metadata_file).load()
-        if getattr(metadata, "storage_meta", None) is None:
-            metadata.storage_meta = dist_cp.StorageMeta()
-        metadata.storage_meta.load_id = self.load_id
-        if metadata.planner_data is None:
-            metadata.planner_data = {}
-        return metadata
-
-    def _collect_read_operations(self, plan):
-        read_ops = []
-
-        def collect_from_storage_metadata(storage_item, base_path=""):
-            if isinstance(storage_item, dist_cp.metadata.BytesStorageMetadata):
-                full_path = self.fs.concat_path(self.path, storage_item.relative_path)
-                read_ops.append(("bytes", storage_item.relative_path, full_path, None, None))
-            elif isinstance(storage_item, dist_cp.metadata.ChunkStorageMetadata):
-                full_path = self.fs.concat_path(self.path, storage_item.relative_path)
-                read_ops.append(
-                    ("chunk", storage_item.relative_path, full_path, storage_item.offset, storage_item.length)
-                )
-            elif isinstance(storage_item, dist_cp.metadata.TensorStorageMetadata):
-                if hasattr(storage_item, "chunks"):
-                    for chunk in storage_item.chunks:
-                        collect_from_storage_metadata(chunk)
-
-        if hasattr(plan, "items"):
-            for item in plan.items:
-                collect_from_storage_metadata(item)
-        elif hasattr(plan, "state_dict_metadata"):
-            for _key, metadata in plan.state_dict_metadata.items():
-                if hasattr(metadata, "storage"):
-                    collect_from_storage_metadata(metadata.storage)
-
-        return read_ops
-
-    def _parallel_preload_files(self, metadata=None, keys=None, state_dict_metadata_dict=None):
-        if not self.max_workers or self.max_workers <= 1:
-            return None
-
-        keys_set = set(keys) if keys else None
-
-        read_ops = []
-
-        if state_dict_metadata_dict is not None:
-            metadata_dict = state_dict_metadata_dict
-        elif metadata is not None and hasattr(metadata, "state_dict_metadata"):
-            metadata_dict = metadata.state_dict_metadata
-        else:
-            if metadata is None:
-                metadata = self.read_metadata()
-            metadata_dict = metadata.state_dict_metadata if hasattr(metadata, "state_dict_metadata") else {}
-
-        for key, tensor_meta in metadata_dict.items():
-            if "optimizer" in key or "_state" in key:
-                continue
-            if keys_set is not None and key not in keys_set:
-                continue
-
-            if hasattr(tensor_meta, "storage"):
-                storage = tensor_meta.storage
-                if isinstance(storage, dist_cp.metadata.BytesStorageMetadata):
-                    full_path = self.fs.concat_path(self.path, storage.relative_path)
-                    read_ops.append(("bytes", storage.relative_path, full_path, None, None))
-                elif isinstance(storage, dist_cp.metadata.ChunkStorageMetadata):
-                    full_path = self.fs.concat_path(self.path, storage.relative_path)
-                    read_ops.append(("chunk", storage.relative_path, full_path, storage.offset, storage.length))
-                elif hasattr(storage, "chunks"):
-                    for chunk in storage.chunks:
-                        if isinstance(chunk, dist_cp.metadata.ChunkStorageMetadata):
-                            full_path = self.fs.concat_path(self.path, chunk.relative_path)
-                            read_ops.append(("chunk", chunk.relative_path, full_path, chunk.offset, chunk.length))
-
-        if not read_ops:
-            return None
-
-        # Group by file path - for chunked files, we'll cache the whole file
-        file_ops = {}
-        for op_type, rel_path, _full_path, offset, length in read_ops:
-            if rel_path not in file_ops:
-                file_ops[rel_path] = []
-            file_ops[rel_path].append((op_type, offset, length))
-
-        # Parallel read files - for files with chunks, read the whole file
-        def read_file(rel_path, full_path, ops):
-            try:
-                with self.fs.create_stream(full_path, "rb") as stream:
-                    has_chunks = any(op[0] == "chunk" and op[1] is not None for op in ops)
-                    if has_chunks:
-                        data = stream.read()
-                        return rel_path, data, True  # True indicates full file
-                    else:
-                        data = stream.read()
-                        return rel_path, data, False
-            except Exception as e:
-                print(f"Error reading {rel_path}: {e}")
-                return None
-
-        print(f"Parallel pre-loading {len(file_ops)} files with {self.max_workers} workers...")
-        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
-            futures = {
-                executor.submit(read_file, rel_path, full_path, ops): rel_path
-                for rel_path, (full_path, ops) in [
-                    (r, (self.fs.concat_path(self.path, r), ops)) for r, ops in file_ops.items()
-                ]
-            }
-            for future in tqdm(as_completed(futures), total=len(futures), desc="Pre-loading files"):
-                result = future.result()
-                if result:
-                    rel_path, data, is_full_file = result
-                    with self._cache_lock:
-                        # Cache by file path - for chunked files, store the full file
-                        self._read_cache[rel_path] = data
-
-        return self._read_cache
-
-    @override
-    def read_data(self, plan, planner):
-        """Override read_data to use parallel reading when enabled"""
-        if self.max_workers and self.max_workers > 1 and self._read_cache is not None and len(self._read_cache) > 0:
-            original_create_stream = self.fs.create_stream
-
-            def cached_create_stream(path, mode="rb"):
-                if mode == "rb":
-                    try:
-                        if path.startswith(self.path):
-                            rel_path = os.path.relpath(path, self.path)
-                        else:
-                            filename = os.path.basename(path)
-                            rel_path = None
-                            for cached_path in self._read_cache.keys():
-                                if os.path.basename(cached_path) == filename:
-                                    rel_path = cached_path
-                                    break
-                            if rel_path is None:
-                                return original_create_stream(path, mode)
-
-                        if rel_path in self._read_cache:
-                            from io import BytesIO
-
-                            return BytesIO(self._read_cache[rel_path])
-                    except (ValueError, KeyError, OSError):
-                        pass
-                return original_create_stream(path, mode)
-
-            self.fs.create_stream = cached_create_stream
-            try:
-                result = super().read_data(plan, planner)
-            finally:
-                self.fs.create_stream = original_create_stream
-            return result
-
-        return super().read_data(plan, planner)
-
-
-class EmptyStateDictLoadPlanner(dist_cp.default_planner.DefaultLoadPlanner):
-    def __init__(self, keys=None):
-        super().__init__()
-        self.keys = set(keys) if keys else None
-
-    @override
-    def set_up_planner(
-        self,
-        state_dict: dist_cp.metadata.STATE_DICT_TYPE,
-        metadata: dist_cp.metadata.Metadata | None = None,
-        is_coordinator: bool = False,
-    ) -> None:
-        for k, v in metadata.state_dict_metadata.items():
-            if "optimizer" in k or "_state" in k:
-                continue
-            if self.keys is not None and k not in self.keys:
-                continue
-            print(f"find {k} in torch_dist ckpt")
-            if isinstance(v, dist_cp.metadata.TensorStorageMetadata):
-                v = torch.empty(v.size, dtype=v.properties.dtype)  # type: ignore[assignment]
-            state_dict[k] = v
-        super().set_up_planner(state_dict, metadata, is_coordinator)
-
-
-def get_expert_param(args, name, param):
-    if ".experts." not in name:
-        yield name, param
-        return
-
-    num_experts = args.num_experts
-    match = re.search(r"mlp.experts\.(.+)\.weight(\d+)", name)
-    if not match:
-        assert param.shape[0] == num_experts
-        for expert_id in range(num_experts):
-            expert_name = name.replace(".experts.experts.", ".experts.") + str(expert_id)
-            expert_param = param[expert_id]
-            yield expert_name, expert_param
-    else:
-        yield name, param
-
-
-def get_layer_param(args, name, param):
-    if ".layers." not in name:
-        yield name, param
-        return
-
-    num_layers = args.num_layers
-    match = re.search(r"\.layers\.(\d+)\.", name)
-    if not match:
-        assert param.shape[0] == num_layers
-        for layer_id in range(num_layers):
-            layer_name = name.replace(".layers.", f".layers.{layer_id}.")
-            layer_param = param[layer_id]
-            yield from get_expert_param(args, layer_name, layer_param)
-    else:
-        yield from get_expert_param(args, name, param)
-
-
-def get_named_params(args, state_dict):
-    for name, param in state_dict.items():
-        name = f"module.module.{name}"
-        yield from get_layer_param(args, name, param)
-
-
-def process_param(args, model_name, name, param, vocab_size=None):
-    if vocab_size:
-        param = remove_padding(name, param, vocab_size)
-    converted_named_tensors = list(convert_to_hf(args, model_name, name, param))
-    return converted_named_tensors
-
-
-def save_tensors(args, model_name, state_dict, output_dir, chunk_size, vocab_size=None, max_workers=1, worker_id=None):
-    # for slime update_weight compatible
-    args.sglang_enable_ep_moe = False
-
-    print(f"start saving to {output_dir}")
-    os.makedirs(output_dir, exist_ok=True)
-    param_list = list(get_named_params(args, state_dict))
-    print(f"Total parameters to process: {len(param_list)}")
-
-    all_converted_tensors = []
-    lock = threading.Lock()
-
-    def process_and_collect(name_param_pair):
-        name, param = name_param_pair
-        try:
-            converted = process_param(args, model_name, name, param, vocab_size)
-            return converted
-        except Exception as e:
-            print(f"Error processing {name}: {e}")
-            return []
-
-    print(f"Processing with {max_workers} workers")
-    with ThreadPoolExecutor(max_workers=max_workers) as executor:
-        futures = {executor.submit(process_and_collect, (name, param)): name for name, param in param_list}
-        for future in tqdm(as_completed(futures), total=len(futures), desc="Converting parameters"):
-            converted = future.result()
-            with lock:
-                all_converted_tensors.extend(converted)
-
-    current_size = 0
-    total_size = 0
-    modeltensors = [{}]
-    for converted_name, converted_param in all_converted_tensors:
-        tensor_size = converted_param.numel() * converted_param.element_size()
-        if tensor_size + current_size > chunk_size:
-            modeltensors.append({})
-            current_size = 0
-        modeltensors[-1][converted_name] = converted_param
-        current_size += tensor_size
-        total_size += tensor_size
-
-    metadata = {"metadata": {"total_size": total_size}, "weight_map": {}}
-
-    num_files = len(modeltensors)
-    file_prefix = f"worker_{worker_id}_" if worker_id is not None else ""
-    index_prefix = f"worker_{worker_id}_" if worker_id is not None else ""
-
-    for i, tensors in enumerate(modeltensors):
-        filename = f"{file_prefix}model-{i:05d}-of-{num_files:05d}.safetensors"
-        for key in tensors.keys():
-            metadata["weight_map"][key] = filename
-    index_filepath = os.path.join(output_dir, f"{index_prefix}model.safetensors.index.json")
-    json.dump(metadata, open(index_filepath, "w"), indent=2)
-    print(f"{index_filepath} saved.")
-
-    def save_file(i, tensors):
-        filename = f"{file_prefix}model-{i:05d}-of-{num_files:05d}.safetensors"
-        t = time.time()
-        filepath = os.path.join(output_dir, filename)
-        safetensors.torch.save_file(tensors, filepath)
-        return filename, time.time() - t
-
-    with ThreadPoolExecutor(max_workers=min(max_workers, num_files)) as executor:
-        futures = [executor.submit(save_file, i, tensors) for i, tensors in enumerate(modeltensors)]
-        for future in tqdm(as_completed(futures), total=len(futures), desc="Saving files"):
-            filename, elapsed = future.result()
-            print(f"{filename} saved in {elapsed:.2f} sec.")
-
-
-def copy_assets(origin_hf_dir, output_dir):
-    for filename in os.listdir(origin_hf_dir):
-        if filename == "model.safetensors.index.json" or filename.endswith(".safetensors"):
-            continue
-        origin_filename = os.path.join(origin_hf_dir, filename)
-        if not os.path.isfile(origin_filename):
-            print(f"Skip {filename}, not a file.")
-            continue
-        src, dst = origin_filename, os.path.join(output_dir, filename)
-        print(f"copy from {src} to {dst}")
-        shutil.copy(src, dst)
-
-
-def conversion_worker(
-    worker_id,
-    keys,
-    args,
-    megatron_args,
-    model_name,
-    state_dict_metadata_dict=None,
-    load_max_workers=2,
-    save_max_workers=16,
-):
-    print(f"Worker {worker_id} starting with {len(keys)} keys")
-    planner = EmptyStateDictLoadPlanner(keys)
-    state_dict = {}
-
-    storage_reader = WrappedStorageReader(args.input_dir, max_workers=load_max_workers)
-
-    if load_max_workers and load_max_workers > 1:
-        print(f"Worker {worker_id}: Pre-loading files with {load_max_workers} workers...")
-        storage_reader._parallel_preload_files(keys=keys, state_dict_metadata_dict=state_dict_metadata_dict)
-
-    print(f"Worker {worker_id}: Loading state dict...")
-    t = time.time()
-    dist_cp.state_dict_loader._load_state_dict(
-        state_dict,
-        storage_reader=storage_reader,
-        planner=planner,
-        no_dist=True,
-    )
-    print(f"Worker {worker_id}: State dict loaded in {time.time()-t:.2f} sec.")
-
-    save_tensors(
-        megatron_args,
-        model_name,
-        state_dict,
-        args.output_dir,
-        args.chunk_size,
-        args.vocab_size,
-        max_workers=save_max_workers,
-        worker_id=worker_id,
-    )
-
-    index_file = os.path.join(args.output_dir, f"worker_{worker_id}_model.safetensors.index.json")
-    if os.path.exists(index_file):
-        with open(index_file) as f:
-            idx = json.load(f)
-        return idx
-    return None
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model-name", type=str, default=None)
-    parser.add_argument("--input-dir", type=str, required=True)
-    parser.add_argument("--output-dir", type=str, required=True)
-    parser.add_argument(
-        "--origin-hf-dir",
-        type=str,
-        default=None,
-        help="use the origin hf dir to copy files like tokenizer, config.json, etc.",
-    )
-    parser.add_argument(
-        "-f", "--force", action="store_true", help="Force overwrite the output directory if it exists."
-    )
-    parser.add_argument(
-        "--chunk-size",
-        type=int,
-        default=5 * 1024**3,
-        help="Chunk size for saving tensors, default is 2GB.",
-    )
-    parser.add_argument(
-        "--vocab-size",
-        type=int,
-        default=None,
-        help="Vocab size for removing padding, if applicable. If not provided, no padding will be removed.",
-    )
-    parser.add_argument(
-        "--load-max-workers",
-        type=int,
-        default=2,
-        help="Number of worker threads for parallel file loading. Default is 2.",
-    )
-    parser.add_argument(
-        "--save-max-workers",
-        type=int,
-        default=16,
-        help="Number of worker threads for parallel tensor conversion and saving. Default is 16.",
-    )
-    args = parser.parse_args()
-
-    if os.path.exists(args.output_dir) and not args.force:
-        raise ValueError(f"Output directory {args.output_dir} already exists. Use --force to overwrite it.")
-
-    if args.model_name is None and args.origin_hf_dir is None:
-        raise ValueError(
-            "Either --model-name or --origin-hf-dir must be provided, so that we can know the name of the params."
-        )
-
-    if args.model_name is None:
-        hf_config = AutoConfig.from_pretrained(args.origin_hf_dir, trust_remote_code=True)
-        args.model_name = type(hf_config).__name__.lower()
-
-    megatron_args = torch.load(os.path.join(args.input_dir, "common.pt"), weights_only=False)["args"]
-
-    load_max_workers = args.load_max_workers
-    save_max_workers = args.save_max_workers
-
-    reader = WrappedStorageReader(args.input_dir)
-    metadata = reader.read_metadata()
-    all_keys = [k for k in metadata.state_dict_metadata.keys() if "optimizer" not in k and "_state" not in k]
-
-    # Group paired keys (linear_q_down_proj and linear_kv_down_proj) together
-    # These will be converted to q_a_proj and kv_a_proj_with_mqa later
-    paired_keys = set()
-    key_groups = []  # Each element is either a single key or a pair of keys
-
-    for key in all_keys:
-        if key in paired_keys:
-            continue
-
-        # Check if this is a linear_q_down_proj or linear_kv_down_proj key
-        if "linear_q_down_proj" in key or "linear_kv_down_proj" in key:
-            # Find its pair
-            if "linear_q_down_proj" in key:
-                pair_key = key.replace("linear_q_down_proj", "linear_kv_down_proj")
-            else:
-                pair_key = key.replace("linear_kv_down_proj", "linear_q_down_proj")
-
-            # If pair exists in all_keys, group them together
-            if pair_key in all_keys:
-                key_groups.append([key, pair_key])
-                paired_keys.add(key)
-                paired_keys.add(pair_key)
-            else:
-                key_groups.append([key])
-        else:
-            key_groups.append([key])
-
-    # Distribute key groups to workers (round-robin to balance load)
-    num_workers = load_max_workers
-    key_chunks = [[] for _ in range(num_workers)]
-
-    for i, key_group in enumerate(key_groups):
-        worker_idx = i % num_workers
-        key_chunks[worker_idx].extend(key_group)
-
-    # Remove empty chunks
-    key_chunks = [chunk for chunk in key_chunks if chunk]
-    num_workers = len(key_chunks)
-
-    print(
-        f"Total keys: {len(all_keys)}, Paired groups: {len([g for g in key_groups if len(g) > 1])}, Workers: {num_workers}"
-    )
-
-    state_dict_metadata_dict = metadata.state_dict_metadata if hasattr(metadata, "state_dict_metadata") else {}
-
-    ctx = multiprocessing.get_context("spawn")
-    with ctx.Pool(num_workers) as pool:
-        results = pool.starmap(
-            conversion_worker,
-            [
-                (
-                    i,
-                    key_chunks[i],
-                    args,
-                    megatron_args,
-                    args.model_name,
-                    state_dict_metadata_dict,
-                    load_max_workers,
-                    save_max_workers,
-                )
-                for i in range(num_workers)
-            ],
-        )
-
-    final_weight_map = {}
-    total_size = 0
-    final_file_index = 1
-
-    for i, res in enumerate(results):
-        if not res:
-            continue
-
-        w_map = res["weight_map"]
-        total_size += res["metadata"]["total_size"]
-
-        worker_files = sorted(list(set(w_map.values())))
-        file_rename_map = {}
-        for wf in worker_files:
-            new_name = f"model-{final_file_index:05d}.safetensors"
-            final_file_index += 1
-            file_rename_map[wf] = new_name
-
-            src = os.path.join(args.output_dir, wf)
-            dst = os.path.join(args.output_dir, new_name)
-            if os.path.exists(src):
-                shutil.move(src, dst)
-
-        for param_name, old_fname in w_map.items():
-            final_weight_map[param_name] = file_rename_map[old_fname]
-
-        temp_index = os.path.join(args.output_dir, f"worker_{i}_model.safetensors.index.json")
-        if os.path.exists(temp_index):
-            os.remove(temp_index)
-
-    total_files = final_file_index - 1
-    final_weight_map_fixed = {}
-    for i in range(1, total_files + 1):
-        old_name = f"model-{i:05d}.safetensors"
-        new_name = f"model-{i:05d}-of-{total_files:05d}.safetensors"
-        old_path = os.path.join(args.output_dir, old_name)
-        new_path = os.path.join(args.output_dir, new_name)
-        if os.path.exists(old_path):
-            shutil.move(old_path, new_path)
-        for k, v in final_weight_map.items():
-            if v == old_name:
-                final_weight_map_fixed[k] = new_name
-
-    index_data = {"metadata": {"total_size": total_size}, "weight_map": final_weight_map_fixed}
-    json.dump(index_data, open(os.path.join(args.output_dir, "model.safetensors.index.json"), "w"), indent=2)
-    print("Model converted and saved.")
-
-    if args.origin_hf_dir:
-        copy_assets(args.origin_hf_dir, args.output_dir)
diff --git a/tools/replay_openai_jsonl.py b/tools/replay_openai_jsonl.py
deleted file mode 100644
index fd7b4a2f94..0000000000
--- a/tools/replay_openai_jsonl.py
+++ /dev/null
@@ -1,535 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import asyncio
-import json
-import math
-import time
-from dataclasses import dataclass, field
-from pathlib import Path
-from typing import Any
-
-import httpx
-
-DEFAULT_INPUT_FILE = Path(__file__).resolve().parent.parent / "tob-glm-5.filtered.jsonl"
-STREAM_MODE_CHOICES = ("inherit", "true", "false")
-
-
-@dataclass
-class ReplayStats:
-    submitted: int = 0
-    completed: int = 0
-    succeeded: int = 0
-    failed: int = 0
-    stream_requests: int = 0
-    status_counts: dict[str, int] = field(default_factory=dict)
-    latencies_ms: list[float] = field(default_factory=list)
-
-    def record(self, result: dict[str, Any]) -> None:
-        self.completed += 1
-        if result["ok"]:
-            self.succeeded += 1
-        else:
-            self.failed += 1
-
-        if result.get("stream"):
-            self.stream_requests += 1
-
-        status_key = str(result.get("status_code", "transport_error"))
-        self.status_counts[status_key] = self.status_counts.get(status_key, 0) + 1
-
-        latency_ms = result.get("latency_ms")
-        if isinstance(latency_ms, (int, float)):
-            self.latencies_ms.append(float(latency_ms))
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Replay OpenAI-compatible chat completion payloads from a JSONL file against an SGLang router.",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-        epilog=(
-            "Examples:\n"
-            "  python tools/replay_openai_jsonl.py --port 30000 --concurrency 32\n"
-            "  python tools/replay_openai_jsonl.py --port 30000 --replay-times 1\n"
-            "  python tools/replay_openai_jsonl.py --base-url http://127.0.0.1:30000/v1 --concurrency 64 --max-requests 100\n"
-            "  python tools/replay_openai_jsonl.py --port 30000 --model custom --stream-mode false --output-file /tmp/replay.jsonl"
-        ),
-    )
-    parser.add_argument("--input-file", type=Path, default=DEFAULT_INPUT_FILE, help="Input JSONL file to replay")
-    parser.add_argument("--host", default="127.0.0.1", help="Router host when --base-url is not provided")
-    parser.add_argument("--port", type=int, help="Router port when --base-url is not provided")
-    parser.add_argument(
-        "--base-url",
-        help="OpenAI-compatible base URL, e.g. http://127.0.0.1:30000/v1. If omitted, it is built from --host/--port.",
-    )
-    parser.add_argument("--endpoint", default="/chat/completions", help="Endpoint path to append to the base URL")
-    parser.add_argument("--api-key", help="Optional Bearer token for the OpenAI-compatible endpoint")
-    parser.add_argument("--concurrency", type=int, default=32, help="Maximum number of in-flight requests")
-    parser.add_argument("--timeout", type=float, default=600.0, help="Per-request timeout in seconds")
-    parser.add_argument("--retries", type=int, default=0, help="Retry count for 429/5xx/transport failures")
-    parser.add_argument("--max-requests", type=int, help="Stop after replaying at most this many input lines")
-    parser.add_argument("--replay-times", type=int, default=10, help="Replay the input dataset this many times")
-    parser.add_argument("--model", help="Override model for every request")
-    parser.add_argument("--default-model", default="custom", help="Fallback model when the payload has no model")
-    parser.add_argument(
-        "--stream-mode",
-        choices=STREAM_MODE_CHOICES,
-        default="inherit",
-        help="Whether to keep the payload stream flag or override it",
-    )
-    parser.add_argument("--output-file", type=Path, help="Optional JSONL file for per-request results")
-    parser.add_argument("--overwrite", action="store_true", help="Overwrite --output-file if it already exists")
-    parser.add_argument(
-        "--progress-every",
-        type=int,
-        default=100,
-        help="Print one progress line every N completed requests; set <= 0 to disable",
-    )
-    args = parser.parse_args()
-
-    if not args.base_url and args.port is None:
-        parser.error("Either --base-url or --port must be provided.")
-    if args.concurrency <= 0:
-        parser.error("--concurrency must be positive.")
-    if args.timeout <= 0:
-        parser.error("--timeout must be positive.")
-    if args.retries < 0:
-        parser.error("--retries cannot be negative.")
-    if args.max_requests is not None and args.max_requests <= 0:
-        parser.error("--max-requests must be positive when provided.")
-    if args.replay_times <= 0:
-        parser.error("--replay-times must be positive.")
-
-    return args
-
-
-def build_request_url(args: argparse.Namespace) -> str:
-    endpoint = args.endpoint if args.endpoint.startswith("/") else f"/{args.endpoint}"
-    if args.base_url:
-        base_url = args.base_url.rstrip("/")
-    else:
-        base_url = f"http://{args.host}:{args.port}/v1"
-    if base_url.endswith(endpoint):
-        return base_url
-    return f"{base_url}{endpoint}"
-
-
-def build_headers(api_key: str | None, stream: bool) -> dict[str, str]:
-    headers = {
-        "Content-Type": "application/json",
-        "Accept": "text/event-stream" if stream else "application/json",
-    }
-    if api_key:
-        headers["Authorization"] = f"Bearer {api_key}"
-    return headers
-
-
-def normalize_payload(raw_payload: dict[str, Any], args: argparse.Namespace) -> dict[str, Any]:
-    payload = dict(raw_payload)
-
-    if args.model is not None:
-        payload["model"] = args.model
-    elif not str(payload.get("model", "") or "").strip():
-        payload["model"] = args.default_model
-
-    if args.stream_mode != "inherit":
-        payload["stream"] = args.stream_mode == "true"
-        if not payload["stream"]:
-            payload.pop("stream_options", None)
-
-    return payload
-
-
-def decode_body(content: bytes) -> Any:
-    try:
-        return json.loads(content)
-    except json.JSONDecodeError:
-        return content.decode("utf-8", errors="replace")
-
-
-def extract_delta_text(delta: Any, key: str) -> list[str]:
-    if not isinstance(delta, dict):
-        return []
-
-    value = delta.get(key)
-    if isinstance(value, str):
-        return [value]
-
-    if isinstance(value, list):
-        text_parts = []
-        for item in value:
-            if isinstance(item, dict) and isinstance(item.get("text"), str):
-                text_parts.append(item["text"])
-        return text_parts
-
-    return []
-
-
-async def dispatch_request(
-    client: httpx.AsyncClient,
-    request_url: str,
-    api_key: str | None,
-    payload: dict[str, Any],
-) -> tuple[int, str | None, Any]:
-    stream = bool(payload.get("stream", False))
-    headers = build_headers(api_key, stream=stream)
-
-    if not stream:
-        response = await client.post(request_url, json=payload, headers=headers)
-        body = decode_body(response.content)
-        return response.status_code, response.headers.get("x-request-id"), body
-
-    async with client.stream("POST", request_url, json=payload, headers=headers) as response:
-        request_id = response.headers.get("x-request-id")
-        content_type = response.headers.get("content-type", "")
-        if "text/event-stream" not in content_type:
-            body = decode_body(await response.aread())
-            return response.status_code, request_id, body
-
-        content_parts = []
-        reasoning_parts = []
-        finish_reasons = []
-        tool_call_deltas = []
-        usage = None
-        chunk_count = 0
-        raw_events = 0
-
-        async for line in response.aiter_lines():
-            if not line or not line.startswith("data:"):
-                continue
-
-            data = line[5:].strip()
-            if data == "[DONE]":
-                continue
-
-            chunk_count += 1
-            try:
-                event = json.loads(data)
-            except json.JSONDecodeError:
-                raw_events += 1
-                continue
-
-            event_usage = event.get("usage")
-            if event_usage is not None:
-                usage = event_usage
-
-            for choice in event.get("choices", []):
-                finish_reason = choice.get("finish_reason")
-                if finish_reason is not None:
-                    finish_reasons.append(finish_reason)
-
-                delta = choice.get("delta", {})
-                content_parts.extend(extract_delta_text(delta, "content"))
-                reasoning_parts.extend(extract_delta_text(delta, "reasoning"))
-                reasoning_parts.extend(extract_delta_text(delta, "reasoning_content"))
-
-                if delta.get("tool_calls") is not None:
-                    tool_call_deltas.append(delta["tool_calls"])
-
-        body = {
-            "mode": "stream",
-            "content": "".join(content_parts),
-            "reasoning": "".join(reasoning_parts) or None,
-            "finish_reasons": finish_reasons or None,
-            "usage": usage,
-            "chunk_count": chunk_count,
-        }
-        if tool_call_deltas:
-            body["tool_call_deltas"] = tool_call_deltas
-        if raw_events:
-            body["raw_event_count"] = raw_events
-        return response.status_code, request_id, body
-
-
-def should_retry_status(status_code: int) -> bool:
-    return status_code == 429 or status_code >= 500
-
-
-def percentile(values: list[float], ratio: float) -> float | None:
-    if not values:
-        return None
-
-    ordered = sorted(values)
-    if len(ordered) == 1:
-        return ordered[0]
-
-    index = (len(ordered) - 1) * ratio
-    lower = math.floor(index)
-    upper = math.ceil(index)
-    if lower == upper:
-        return ordered[lower]
-
-    lower_value = ordered[lower]
-    upper_value = ordered[upper]
-    return lower_value + (upper_value - lower_value) * (index - lower)
-
-
-def make_summary(
-    stats: ReplayStats,
-    elapsed_s: float,
-    request_url: str,
-    output_file: Path | None,
-    replay_times: int,
-) -> dict[str, Any]:
-    throughput = stats.completed / elapsed_s if elapsed_s > 0 else None
-    return {
-        "request_url": request_url,
-        "replay_times": replay_times,
-        "submitted": stats.submitted,
-        "completed": stats.completed,
-        "succeeded": stats.succeeded,
-        "failed": stats.failed,
-        "stream_requests": stats.stream_requests,
-        "non_stream_requests": stats.completed - stats.stream_requests,
-        "status_counts": stats.status_counts,
-        "elapsed_s": round(elapsed_s, 3),
-        "throughput_rps": round(throughput, 3) if throughput is not None else None,
-        "latency_ms": {
-            "p50": percentile(stats.latencies_ms, 0.50),
-            "p95": percentile(stats.latencies_ms, 0.95),
-            "p99": percentile(stats.latencies_ms, 0.99),
-        },
-        "output_file": str(output_file) if output_file else None,
-    }
-
-
-def print_progress(stats: ReplayStats) -> None:
-    print(
-        f"[progress] completed={stats.completed} succeeded={stats.succeeded} failed={stats.failed} "
-        f"stream={stats.stream_requests}"
-    )
-
-
-def emit_result(result: dict[str, Any], output_handle) -> None:
-    if output_handle is None:
-        return
-    output_handle.write(json.dumps(result, ensure_ascii=False) + "\n")
-
-
-def validate_output_path(output_file: Path | None, overwrite: bool) -> None:
-    if output_file is None:
-        return
-    if output_file.exists() and not overwrite:
-        raise SystemExit(f"Output file already exists: {output_file}. Pass --overwrite to replace it.")
-    output_file.parent.mkdir(parents=True, exist_ok=True)
-
-
-def make_input_error_result(replay_round: int, line_number: int, error: str) -> dict[str, Any]:
-    return {
-        "replay_round": replay_round,
-        "line_number": line_number,
-        "ok": False,
-        "status_code": None,
-        "request_id": None,
-        "latency_ms": 0.0,
-        "stream": None,
-        "model": None,
-        "attempts": 0,
-        "error": error,
-        "response": None,
-    }
-
-
-async def send_request(
-    client: httpx.AsyncClient,
-    request_url: str,
-    api_key: str | None,
-    replay_round: int,
-    line_number: int,
-    payload: dict[str, Any],
-    retries: int,
-) -> dict[str, Any]:
-    started = time.perf_counter()
-    last_status_code = None
-    last_request_id = None
-    last_response = None
-    last_error = None
-
-    for attempt in range(1, retries + 2):
-        try:
-            status_code, request_id, response_body = await dispatch_request(client, request_url, api_key, payload)
-        except httpx.HTTPError as exc:
-            last_error = f"{type(exc).__name__}: {exc}"
-            if attempt <= retries:
-                await asyncio.sleep(min(0.5 * attempt, 2.0))
-                continue
-            latency_ms = round((time.perf_counter() - started) * 1000, 3)
-            return {
-                "replay_round": replay_round,
-                "line_number": line_number,
-                "ok": False,
-                "status_code": None,
-                "request_id": None,
-                "latency_ms": latency_ms,
-                "stream": bool(payload.get("stream", False)),
-                "model": payload.get("model"),
-                "attempts": attempt,
-                "error": last_error,
-                "response": None,
-            }
-
-        last_status_code = status_code
-        last_request_id = request_id
-        last_response = response_body
-
-        if status_code < 400:
-            latency_ms = round((time.perf_counter() - started) * 1000, 3)
-            return {
-                "replay_round": replay_round,
-                "line_number": line_number,
-                "ok": True,
-                "status_code": status_code,
-                "request_id": request_id,
-                "latency_ms": latency_ms,
-                "stream": bool(payload.get("stream", False)),
-                "model": payload.get("model"),
-                "attempts": attempt,
-                "error": None,
-                "response": response_body,
-            }
-
-        last_error = f"HTTP {status_code}"
-        if attempt <= retries and should_retry_status(status_code):
-            await asyncio.sleep(min(0.5 * attempt, 2.0))
-            continue
-
-        break
-
-    latency_ms = round((time.perf_counter() - started) * 1000, 3)
-    return {
-        "replay_round": replay_round,
-        "line_number": line_number,
-        "ok": False,
-        "status_code": last_status_code,
-        "request_id": last_request_id,
-        "latency_ms": latency_ms,
-        "stream": bool(payload.get("stream", False)),
-        "model": payload.get("model"),
-        "attempts": retries + 1,
-        "error": last_error,
-        "response": last_response,
-    }
-
-
-async def collect_completed(
-    pending: set[asyncio.Task],
-    stats: ReplayStats,
-    output_handle,
-    progress_every: int,
-) -> set[asyncio.Task]:
-    done, still_pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
-    for task in done:
-        result = await task
-        emit_result(result, output_handle)
-        stats.record(result)
-        if progress_every > 0 and stats.completed % progress_every == 0:
-            print_progress(stats)
-    return still_pending
-
-
-async def run(args: argparse.Namespace) -> int:
-    validate_output_path(args.output_file, overwrite=args.overwrite)
-    request_url = build_request_url(args)
-    input_file = args.input_file.resolve()
-    if not input_file.exists():
-        raise SystemExit(f"Input file does not exist: {input_file}")
-
-    print(
-        f"Replaying {input_file} -> {request_url} with concurrency={args.concurrency}, replay_times={args.replay_times}"
-    )
-
-    stats = ReplayStats()
-    pending: set[asyncio.Task] = set()
-    output_handle = args.output_file.open("w", encoding="utf-8") if args.output_file else None
-    started = time.perf_counter()
-
-    try:
-        limits = httpx.Limits(
-            max_connections=args.concurrency,
-            max_keepalive_connections=max(20, args.concurrency),
-        )
-        timeout = httpx.Timeout(args.timeout)
-        async with httpx.AsyncClient(limits=limits, timeout=timeout, http2=True, trust_env=False) as client:
-            with input_file.open("r", encoding="utf-8") as input_handle:
-                for replay_round in range(1, args.replay_times + 1):
-                    input_handle.seek(0)
-                    for line_number, line in enumerate(input_handle, start=1):
-                        if args.max_requests is not None and stats.submitted >= args.max_requests:
-                            break
-
-                        stripped = line.strip()
-                        if not stripped:
-                            continue
-
-                        try:
-                            raw_payload = json.loads(stripped)
-                        except json.JSONDecodeError as exc:
-                            stats.submitted += 1
-                            result = make_input_error_result(replay_round, line_number, f"Invalid JSON: {exc}")
-                            emit_result(result, output_handle)
-                            stats.record(result)
-                            continue
-
-                        if not isinstance(raw_payload, dict):
-                            stats.submitted += 1
-                            result = make_input_error_result(replay_round, line_number, "JSON line must be an object")
-                            emit_result(result, output_handle)
-                            stats.record(result)
-                            continue
-
-                        payload = normalize_payload(raw_payload, args)
-                        pending.add(
-                            asyncio.create_task(
-                                send_request(
-                                    client=client,
-                                    request_url=request_url,
-                                    api_key=args.api_key,
-                                    replay_round=replay_round,
-                                    line_number=line_number,
-                                    payload=payload,
-                                    retries=args.retries,
-                                )
-                            )
-                        )
-                        stats.submitted += 1
-
-                        if len(pending) >= args.concurrency:
-                            pending = await collect_completed(
-                                pending,
-                                stats=stats,
-                                output_handle=output_handle,
-                                progress_every=args.progress_every,
-                            )
-
-                    if args.max_requests is not None and stats.submitted >= args.max_requests:
-                        break
-
-            while pending:
-                pending = await collect_completed(
-                    pending,
-                    stats=stats,
-                    output_handle=output_handle,
-                    progress_every=args.progress_every,
-                )
-    finally:
-        if output_handle is not None:
-            output_handle.close()
-
-    summary = make_summary(
-        stats,
-        elapsed_s=time.perf_counter() - started,
-        request_url=request_url,
-        output_file=args.output_file,
-        replay_times=args.replay_times,
-    )
-    print(json.dumps(summary, indent=2, ensure_ascii=False))
-    return 0 if stats.failed == 0 else 1
-
-
-def main() -> int:
-    args = parse_args()
-    return asyncio.run(run(args))
-
-
-if __name__ == "__main__":
-    raise SystemExit(main())
diff --git a/tools/trace_timeline_viewer.py b/tools/trace_timeline_viewer.py
deleted file mode 100644
index f08683f2ec..0000000000
--- a/tools/trace_timeline_viewer.py
+++ /dev/null
@@ -1,2370 +0,0 @@
-#!/usr/bin/env python3
-"""Build and serve an interactive timeline viewer for rollout trace dumps.
-
-The viewer consumes a rollout debug dump `.pt` file, extracts per-sample trace
-events, rebuilds spans and point events, and writes a lightweight JSON cache
-plus a self-contained HTML viewer next to the source file.
-"""
-
-from __future__ import annotations
-
-import argparse
-import functools
-import json
-import pickle
-import socketserver
-import sys
-import time
-import types
-from dataclasses import dataclass
-from http.server import SimpleHTTPRequestHandler
-from pathlib import Path
-from typing import Any
-
-import torch
-
-
-CACHE_VERSION = 1
-
-
-class _MissingPickleObject:
-    def __setstate__(self, state: Any) -> None:
-        if isinstance(state, dict):
-            self.__dict__.update(state)
-            return
-        self.__dict__["_raw_state"] = state
-
-
-_MISSING_PICKLE_GLOBALS: set[tuple[str, str]] = set()
-
-
-def _ensure_dummy_module(module_name: str) -> types.ModuleType:
-    module = sys.modules.get(module_name)
-    if isinstance(module, types.ModuleType):
-        return module
-
-    module = types.ModuleType(module_name)
-    sys.modules[module_name] = module
-    if "." in module_name:
-        parent_name, child_name = module_name.rsplit(".", 1)
-        parent = _ensure_dummy_module(parent_name)
-        setattr(parent, child_name, module)
-    return module
-
-
-def _make_dummy_pickle_global(module_name: str, name: str) -> type[_MissingPickleObject]:
-    module = _ensure_dummy_module(module_name)
-    existing = getattr(module, name, None)
-    if isinstance(existing, type):
-        return existing
-
-    dummy_type = type(name, (_MissingPickleObject,), {"__module__": module_name})
-    setattr(module, name, dummy_type)
-    _MISSING_PICKLE_GLOBALS.add((module_name, name))
-    return dummy_type
-
-
-class _DummyFallbackUnpickler(pickle.Unpickler):
-    def find_class(self, module: str, name: str) -> Any:
-        try:
-            return super().find_class(module, name)
-        except (AttributeError, ImportError, ModuleNotFoundError):
-            return _make_dummy_pickle_global(module, name)
-
-
-_DUMMY_FALLBACK_PICKLE_MODULE = types.SimpleNamespace(
-    __name__="pickle",
-    Unpickler=_DummyFallbackUnpickler,
-    load=pickle.load,
-    loads=pickle.loads,
-)
-
-
-@dataclass
-class TimelinePaths:
-    pt_path: Path
-    cache_path: Path
-    html_path: Path
-
-
-def _json_safe(value: Any) -> Any:
-    if isinstance(value, (str, int, float, bool)) or value is None:
-        return value
-    if isinstance(value, dict):
-        return {str(k): _json_safe(v) for k, v in value.items()}
-    if isinstance(value, (list, tuple)):
-        return [_json_safe(v) for v in value]
-    return str(value)
-
-
-def _round_float(value: float | None) -> float | None:
-    if value is None:
-        return None
-    return round(float(value), 6)
-
-
-def _compact_text(value: Any, max_len: int = 256) -> Any:
-    value = _json_safe(value)
-    if not isinstance(value, str):
-        return value
-    if len(value) <= max_len:
-        return value
-    return f"{value[:max_len]}...<truncated:{len(value)}>"
-
-
-def _safe_duration(start: float | None, end: float | None) -> float | None:
-    if start is None or end is None:
-        return None
-    return max(0.0, float(end) - float(start))
-
-
-def _to_sample_dict(sample: Any) -> dict[str, Any]:
-    if hasattr(sample, "to_dict"):
-        sample = sample.to_dict()
-    if isinstance(sample, dict):
-        return sample
-    result = {}
-    for key in (
-        "group_index",
-        "index",
-        "prompt",
-        "response",
-        "response_length",
-        "reward",
-        "metadata",
-        "source",
-        "status",
-        "label",
-        "trace",
-    ):
-        if hasattr(sample, key):
-            result[key] = getattr(sample, key)
-    return result
-
-
-def _infer_source(sample: dict[str, Any], metadata: dict[str, Any]) -> Any:
-    if sample.get("source") not in (None, ""):
-        return sample.get("source")
-    if metadata.get("source") not in (None, ""):
-        return metadata.get("source")
-    if metadata.get("source_name") not in (None, ""):
-        return metadata.get("source_name")
-    for key, value in metadata.items():
-        if "source" in str(key).lower() and value not in (None, ""):
-            return value
-    return None
-
-
-def _event_timestamp(event: dict[str, Any]) -> float | None:
-    ts = event.get("ts")
-    if ts is None:
-        return None
-    try:
-        return float(ts)
-    except (TypeError, ValueError):
-        return None
-
-
-def _normalize_trace_events(trace: dict[str, Any]) -> list[dict[str, Any]]:
-    raw_events = trace.get("events") or []
-    normalized = []
-    active_stack: list[str] = []
-
-    for order, raw_event in enumerate(raw_events):
-        if not isinstance(raw_event, dict):
-            continue
-        ts = _event_timestamp(raw_event)
-        if ts is None:
-            continue
-
-        event = {
-            "order": order,
-            "ts": ts,
-            "type": _json_safe(raw_event.get("type")),
-            "name": _json_safe(raw_event.get("name")),
-            "attempt": int(raw_event.get("attempt", trace.get("attempt", 0)) or 0),
-            "sample_id": _json_safe(raw_event.get("sample_id", trace.get("sample_id"))),
-            "group_id": _json_safe(raw_event.get("group_id", trace.get("group_id"))),
-            "span_id": _json_safe(raw_event.get("span_id")),
-            "parent_span_id": _json_safe(raw_event.get("parent_span_id")),
-            "attrs": _json_safe(raw_event.get("attrs") or {}),
-        }
-        event["inferred_parent_span_id"] = active_stack[-1] if active_stack else None
-        normalized.append(event)
-
-        if event["type"] == "span_start" and event["span_id"]:
-            active_stack.append(event["span_id"])
-            continue
-
-        if event["type"] == "span_end" and event["span_id"]:
-            for idx in range(len(active_stack) - 1, -1, -1):
-                if active_stack[idx] == event["span_id"]:
-                    del active_stack[idx]
-                    break
-
-    return normalized
-
-
-def _span_name(item: dict[str, Any]) -> str:
-    return str(item.get("name") or "span")
-
-
-def _span_type(item: dict[str, Any]) -> str:
-    if item["type"] == "event":
-        return "point_event"
-    if item["type"] == "orphan_end":
-        return "orphan_end"
-    return item["state"]
-
-
-def _compute_span_depths(spans: list[dict[str, Any]]) -> dict[str, int]:
-    span_by_id = {span["span_id"]: span for span in spans if span.get("span_id")}
-    cache: dict[str, int] = {}
-
-    def resolve(span_id: str | None, seen: set[str]) -> int:
-        if not span_id or span_id not in span_by_id:
-            return 0
-        if span_id in cache:
-            return cache[span_id]
-        if span_id in seen:
-            cache[span_id] = 0
-            return 0
-        seen.add(span_id)
-        parent_id = span_by_id[span_id].get("parent_span_id")
-        depth = 0 if not parent_id or parent_id not in span_by_id else resolve(parent_id, seen) + 1
-        cache[span_id] = depth
-        return depth
-
-    for span in spans:
-        span_id = span.get("span_id")
-        if span_id:
-            resolve(span_id, set())
-    return cache
-
-
-def _build_items_from_trace(sample: dict[str, Any], sample_idx: int) -> dict[str, Any] | None:
-    trace = sample.get("trace")
-    if not isinstance(trace, dict):
-        return None
-
-    events = _normalize_trace_events(trace)
-    if not events:
-        return None
-
-    open_starts: dict[str, dict[str, Any]] = {}
-    closed_spans: list[dict[str, Any]] = []
-    point_events: list[dict[str, Any]] = []
-    orphan_ends: list[dict[str, Any]] = []
-    all_timestamps: list[float] = []
-
-    for event in events:
-        all_timestamps.append(event["ts"])
-        event_type = event["type"]
-
-        if event_type == "span_start" and event["span_id"]:
-            open_starts[event["span_id"]] = {
-                "type": "span",
-                "state": "closed_span",
-                "name": event["name"],
-                "start_ts": event["ts"],
-                "end_ts": None,
-                "display_end_ts": None,
-                "attempt": event["attempt"],
-                "span_id": event["span_id"],
-                "parent_span_id": event.get("parent_span_id"),
-                "start_attrs": event.get("attrs") or {},
-                "end_attrs": {},
-            }
-            continue
-
-        if event_type == "span_end":
-            span_id = event.get("span_id")
-            start_record = open_starts.pop(span_id, None) if span_id else None
-            if start_record is None:
-                orphan_ends.append(
-                    {
-                        "type": "orphan_end",
-                        "state": "orphan_end",
-                        "name": event["name"],
-                        "ts": event["ts"],
-                        "attempt": event["attempt"],
-                        "span_id": span_id,
-                        "parent_span_id": event.get("parent_span_id") or event.get("inferred_parent_span_id"),
-                        "attrs": event.get("attrs") or {},
-                    }
-                )
-                continue
-
-            start_record["end_ts"] = event["ts"]
-            start_record["display_end_ts"] = event["ts"]
-            start_record["end_attrs"] = event.get("attrs") or {}
-            closed_spans.append(start_record)
-            continue
-
-        point_events.append(
-            {
-                "type": "event",
-                "state": "point_event",
-                "name": event["name"],
-                "ts": event["ts"],
-                "attempt": event["attempt"],
-                "span_id": None,
-                "parent_span_id": event.get("inferred_parent_span_id"),
-                "attrs": event.get("attrs") or {},
-            }
-        )
-
-    row_end_ts = max(all_timestamps) if all_timestamps else None
-    open_spans = list(open_starts.values())
-    all_spans = closed_spans + open_spans
-    span_depths = _compute_span_depths(all_spans)
-    span_by_id = {span["span_id"]: span for span in all_spans if span.get("span_id")}
-    sibling_groups: dict[str | None, list[dict[str, Any]]] = {}
-
-    for span in all_spans:
-        sibling_groups.setdefault(span.get("parent_span_id"), []).append(span)
-
-    for siblings in sibling_groups.values():
-        siblings.sort(key=lambda item: (item["start_ts"], item["end_ts"] or float("inf"), _span_name(item)))
-
-    def nearest_closed_ancestor_end(span: dict[str, Any]) -> float | None:
-        current_parent = span.get("parent_span_id")
-        while current_parent:
-            parent = span_by_id.get(current_parent)
-            if parent is None:
-                return None
-            if parent.get("end_ts") is not None:
-                return float(parent["end_ts"])
-            current_parent = parent.get("parent_span_id")
-        return None
-
-    for span in open_spans:
-        candidates: list[tuple[float, str]] = []
-        if row_end_ts is not None:
-            candidates.append((row_end_ts, "row_end"))
-
-        siblings = sibling_groups.get(span.get("parent_span_id"), [])
-        for sibling in siblings:
-            if sibling is span:
-                continue
-            sibling_start = float(sibling["start_ts"])
-            if sibling_start > float(span["start_ts"]):
-                candidates.append((sibling_start, "next_sibling_start"))
-                break
-
-        ancestor_end = nearest_closed_ancestor_end(span)
-        if ancestor_end is not None and ancestor_end > float(span["start_ts"]):
-            candidates.append((ancestor_end, "ancestor_end"))
-
-        if candidates:
-            display_end_ts, clipped_by = min(candidates, key=lambda item: item[0])
-            if display_end_ts <= float(span["start_ts"]):
-                display_end_ts = float(span["start_ts"])
-                clipped_by = "self"
-        else:
-            display_end_ts = float(span["start_ts"])
-            clipped_by = "self"
-
-        span["state"] = "open_span"
-        span["display_end_ts"] = display_end_ts
-        span.setdefault("end_attrs", {})
-        span["end_attrs"]["clipped_by"] = clipped_by
-
-    for span in all_spans:
-        span["depth"] = span_depths.get(span.get("span_id") or "", 0)
-        span["lane"] = span["depth"]
-
-    for event in point_events:
-        parent_span_id = event.get("parent_span_id")
-        event["depth"] = span_depths.get(parent_span_id or "", 0)
-        event["lane"] = event["depth"]
-
-    for item in orphan_ends:
-        parent_span_id = item.get("parent_span_id")
-        item["depth"] = span_depths.get(parent_span_id or "", 0)
-        item["lane"] = item["depth"]
-
-    def parent_span_name(parent_span_id: str | None) -> str | None:
-        if not parent_span_id:
-            return None
-        parent = span_by_id.get(parent_span_id)
-        if not parent:
-            return None
-        return parent.get("name")
-
-    all_items: list[dict[str, Any]] = []
-    for span in all_spans:
-        all_items.append(
-            {
-                "type": "span",
-                "state": span["state"],
-                "name": span["name"],
-                "start_ts": _round_float(span["start_ts"]),
-                "end_ts": _round_float(span["end_ts"]),
-                "display_end_ts": _round_float(span["display_end_ts"]),
-                "attempt": span["attempt"],
-                "span_id": span.get("span_id"),
-                "parent_span_id": span.get("parent_span_id"),
-                "parent_span_name": parent_span_name(span.get("parent_span_id")),
-                "lane": span["lane"],
-                "depth": span["depth"],
-                "attrs": {
-                    "start_attrs": _json_safe(span.get("start_attrs") or {}),
-                    "end_attrs": _json_safe(span.get("end_attrs") or {}),
-                },
-            }
-        )
-
-    for event in point_events:
-        all_items.append(
-            {
-                "type": "event",
-                "state": "point_event",
-                "name": event["name"],
-                "ts": _round_float(event["ts"]),
-                "attempt": event["attempt"],
-                "span_id": None,
-                "parent_span_id": event.get("parent_span_id"),
-                "parent_span_name": parent_span_name(event.get("parent_span_id")),
-                "lane": event["lane"],
-                "depth": event["depth"],
-                "attrs": _json_safe(event.get("attrs") or {}),
-            }
-        )
-
-    for item in orphan_ends:
-        all_items.append(
-            {
-                "type": "orphan_end",
-                "state": "orphan_end",
-                "name": item["name"],
-                "ts": _round_float(item["ts"]),
-                "attempt": item["attempt"],
-                "span_id": item.get("span_id"),
-                "parent_span_id": item.get("parent_span_id"),
-                "parent_span_name": parent_span_name(item.get("parent_span_id")),
-                "lane": item["lane"],
-                "depth": item["depth"],
-                "attrs": _json_safe(item.get("attrs") or {}),
-            }
-        )
-
-    pd_lane_specs = [
-        (
-            "prefill",
-            "P",
-            [
-                "pd_prefill_bootstrap_queue_duration",
-                "pd_bootstrap_duration",
-                "pd_alloc_waiting_duration",
-                "pd_prefill_forward_duration",
-                "pd_prefill_transfer_queue_duration",
-            ],
-        ),
-        (
-            "decode",
-            "D",
-            [
-                "pd_decode_prealloc_duration",
-                "pd_decode_transfer_duration",
-                "pd_decode_forward_duration",
-            ],
-        ),
-    ]
-    next_virtual_lane = max((item["lane"] for item in all_items), default=-1)
-    for span in all_spans:
-        if span["state"] != "closed_span" or span.get("end_ts") is None:
-            continue
-        end_attrs = span.get("end_attrs") or {}
-        for role, suffix, keys in pd_lane_specs:
-            role_attrs = {
-                key: value for key in keys if isinstance((value := end_attrs.get(key)), (int, float)) and value > 0
-            }
-            if not role_attrs:
-                continue
-            next_virtual_lane += 1
-            role_attrs.update(
-                {
-                    "timeline_pd_virtual_role": role,
-                    "timeline_pd_parent_name": span["name"],
-                    "timeline_pd_parent_duration": _round_float(_safe_duration(span["start_ts"], span["end_ts"])),
-                }
-            )
-            all_items.append(
-                {
-                    "type": "span",
-                    "state": "closed_span",
-                    "name": f'{span["name"]} [{suffix}]',
-                    "start_ts": _round_float(span["start_ts"]),
-                    "end_ts": _round_float(span["end_ts"]),
-                    "display_end_ts": _round_float(span["display_end_ts"]),
-                    "attempt": span["attempt"],
-                    "span_id": f'{span.get("span_id") or span["name"]}:pd:{role}',
-                    "parent_span_id": span.get("span_id"),
-                    "parent_span_name": span["name"],
-                    "lane": next_virtual_lane,
-                    "depth": next_virtual_lane,
-                    "attrs": {
-                        "start_attrs": {},
-                        "end_attrs": _json_safe(role_attrs),
-                    },
-                }
-            )
-
-    all_items.sort(
-        key=lambda item: (
-            item["lane"],
-            item.get("start_ts", item.get("ts", 0.0)),
-            item.get("display_end_ts", item.get("ts", 0.0)),
-            item["name"],
-        )
-    )
-
-    row_start = min(item.get("start_ts", item.get("ts")) for item in all_items)
-    row_end = max(item.get("display_end_ts", item.get("ts")) for item in all_items)
-    response_lengths = []
-    for item in all_items:
-        attrs = item.get("attrs") or {}
-        for payload in (attrs, attrs.get("start_attrs"), attrs.get("end_attrs")):
-            if not isinstance(payload, dict):
-                continue
-            response_length = payload.get("response_length")
-            if isinstance(response_length, (int, float)):
-                response_lengths.append(int(response_length))
-
-    metadata = sample.get("metadata") or {}
-    if not isinstance(metadata, dict):
-        metadata = {}
-
-    reward = sample.get("reward")
-    if isinstance(reward, dict):
-        reward = _json_safe(reward)
-
-    return {
-        "row_id": sample_idx,
-        "sample_index": sample.get("index", sample_idx),
-        "group_index": sample.get("group_index"),
-        "source": _compact_text(_infer_source(sample, metadata), max_len=64),
-        "status": _compact_text(sample.get("status"), max_len=64),
-        "label": _compact_text(sample.get("label"), max_len=256),
-        "reward": reward,
-        "trace_id": _json_safe(trace.get("trace_id")),
-        "attempt": int(trace.get("attempt", 0) or 0),
-        "start": row_start,
-        "end": row_end,
-        "duration": _round_float(_safe_duration(row_start, row_end)),
-        "lane_count": 1 + max((item["lane"] for item in all_items), default=0),
-        "item_count": len(all_items),
-        "closed_span_count": sum(1 for item in all_items if item["state"] == "closed_span"),
-        "open_span_count": sum(1 for item in all_items if item["state"] == "open_span"),
-        "point_event_count": sum(1 for item in all_items if item["state"] == "point_event"),
-        "orphan_count": sum(1 for item in all_items if item["state"] == "orphan_end"),
-        "total_response_length": sum(response_lengths),
-        "max_response_length": max(response_lengths, default=0),
-        "items": all_items,
-    }
-
-
-def _build_cache_data(pt_path: Path) -> dict[str, Any]:
-    before_missing = len(_MISSING_PICKLE_GLOBALS)
-    data = torch.load(
-        pt_path,
-        map_location="cpu",
-        weights_only=False,
-        pickle_module=_DUMMY_FALLBACK_PICKLE_MODULE,
-    )
-    if len(_MISSING_PICKLE_GLOBALS) > before_missing:
-        missing_names = ", ".join(f"{module}.{name}" for module, name in sorted(_MISSING_PICKLE_GLOBALS))
-        print(
-            f"[trace_timeline_viewer] substituted missing pickle globals with dummy classes: {missing_names}",
-            file=sys.stderr,
-        )
-    samples = data["samples"] if isinstance(data, dict) and "samples" in data else data
-
-    rows: list[dict[str, Any]] = []
-    global_start = None
-    global_end = None
-
-    for sample_idx, raw_sample in enumerate(samples):
-        sample = _to_sample_dict(raw_sample)
-        row = _build_items_from_trace(sample, sample_idx)
-        if row is None:
-            continue
-        rows.append(row)
-        global_start = row["start"] if global_start is None else min(global_start, row["start"])
-        global_end = row["end"] if global_end is None else max(global_end, row["end"])
-
-    return {
-        "cache_version": CACHE_VERSION,
-        "pt_path": str(pt_path),
-        "generated_at": time.time(),
-        "sample_count": len(rows),
-        "global_start": _round_float(global_start),
-        "global_end": _round_float(global_end),
-        "rows": rows,
-    }
-
-
-def _timeline_paths(pt_path: Path) -> TimelinePaths:
-    stem = pt_path.stem
-    directory = pt_path.parent
-    return TimelinePaths(
-        pt_path=pt_path,
-        cache_path=directory / f"{stem}.trace_timeline_cache.json",
-        html_path=directory / f"{stem}.trace_timeline_viewer.html",
-    )
-
-
-def ensure_cache(paths: TimelinePaths, rebuild: bool = False) -> dict[str, Any]:
-    if not rebuild and paths.cache_path.exists() and paths.cache_path.stat().st_mtime >= paths.pt_path.stat().st_mtime:
-        with paths.cache_path.open("r", encoding="utf-8") as handle:
-            cached = json.load(handle)
-        if cached.get("cache_version") == CACHE_VERSION:
-            return cached
-
-    cache_data = _build_cache_data(paths.pt_path)
-    with paths.cache_path.open("w", encoding="utf-8") as handle:
-        json.dump(cache_data, handle, ensure_ascii=True, separators=(",", ":"))
-    return cache_data
-
-
-HTML_TEMPLATE = r"""<!doctype html>
-<html lang="en">
-<head>
-  <meta charset="utf-8">
-  <meta name="viewport" content="width=device-width, initial-scale=1">
-  <title>__TITLE__</title>
-  <style>
-    :root {
-      --bg: #f4f1ea;
-      --panel: #fffaf2;
-      --ink: #22201c;
-      --muted: #6f675d;
-      --line: #d8cebf;
-      --accent: #b5502a;
-      --open: #a64f37;
-      --point: #2f6d9b;
-      --orphan: #a12626;
-    }
-    * { box-sizing: border-box; }
-    html, body { height: 100%; overflow: hidden; }
-    body {
-      margin: 0;
-      font-family: "Iosevka Aile", "IBM Plex Sans", sans-serif;
-      color: var(--ink);
-      background: radial-gradient(circle at top left, #fffdf8 0, #f6efe2 34%, #f1e8da 100%);
-    }
-    .page {
-      height: 100vh;
-      overflow: hidden;
-      padding: 18px;
-      display: grid;
-      grid-template-rows: auto auto minmax(0, 1fr);
-      gap: 12px;
-    }
-    .panel {
-      background: rgba(255, 250, 242, 0.92);
-      border: 1px solid var(--line);
-      border-radius: 14px;
-      box-shadow: 0 10px 40px rgba(73, 52, 23, 0.08);
-    }
-    .controls {
-      padding: 14px;
-      display: grid;
-      gap: 10px;
-    }
-    .controls-header {
-      display: flex;
-      gap: 10px;
-      align-items: center;
-      justify-content: space-between;
-    }
-    .controls-extra {
-      display: grid;
-      gap: 10px;
-    }
-    .row {
-      display: flex;
-      gap: 10px;
-      flex-wrap: wrap;
-      align-items: center;
-    }
-    label {
-      display: inline-flex;
-      gap: 6px;
-      align-items: center;
-      color: var(--muted);
-      font-size: 13px;
-    }
-    input, select, button {
-      border: 1px solid var(--line);
-      border-radius: 8px;
-      background: white;
-      padding: 7px 10px;
-      color: var(--ink);
-      font: inherit;
-    }
-    button {
-      cursor: pointer;
-      background: linear-gradient(180deg, #fffdf9, #f4ebde);
-    }
-    .btn-subtle {
-      padding: 6px 9px;
-      border-radius: 8px;
-      background: rgba(255, 255, 255, 0.78);
-      color: #6e6457;
-    }
-    .btn-subtle:hover {
-      background: rgba(255, 255, 255, 0.92);
-      color: #4f463d;
-    }
-    .btn-icon {
-      min-width: 32px;
-      padding: 6px 0;
-      text-align: center;
-      font-size: 14px;
-    }
-    .stats { padding: 14px; }
-    .stats-grid {
-      display: grid;
-      grid-template-columns: repeat(auto-fit, minmax(140px, 1fr));
-      gap: 10px;
-    }
-    .stat {
-      background: rgba(255,255,255,0.66);
-      border: 1px solid var(--line);
-      border-radius: 10px;
-      padding: 10px;
-    }
-    .stat .name {
-      color: var(--muted);
-      font-size: 12px;
-      margin-bottom: 4px;
-      display: flex;
-      align-items: center;
-      gap: 6px;
-    }
-    .stat .value {
-      font-size: 20px;
-      font-weight: 700;
-    }
-    .swatch {
-      width: 12px;
-      height: 8px;
-      border-radius: 2px;
-      display: inline-block;
-      flex: 0 0 auto;
-    }
-    .viewport {
-      overflow: auto;
-      position: relative;
-      padding: 0;
-      min-height: 0;
-      flex: 1 1 auto;
-      overscroll-behavior: contain;
-    }
-    .timeline-panel {
-      min-height: 0;
-      display: flex;
-      flex-direction: column;
-    }
-    .canvas-wrap {
-      position: relative;
-      min-height: 120px;
-    }
-    canvas {
-      display: block;
-      position: sticky;
-      top: 0;
-      left: 0;
-      z-index: 1;
-    }
-    .legend {
-      padding: 0 14px 14px;
-      display: flex;
-      gap: 14px;
-      flex-wrap: wrap;
-      color: var(--muted);
-      font-size: 12px;
-    }
-    .chip {
-      display: inline-flex;
-      align-items: center;
-      gap: 6px;
-    }
-    .dot {
-      width: 10px;
-      height: 10px;
-      border-radius: 999px;
-      display: inline-block;
-    }
-    .bar {
-      width: 14px;
-      height: 8px;
-      border-radius: 2px;
-      display: inline-block;
-    }
-    .line {
-      width: 14px;
-      height: 0;
-      border-top: 2px dashed var(--open);
-      display: inline-block;
-    }
-    .cross {
-      width: 12px;
-      height: 12px;
-      display: inline-block;
-      position: relative;
-    }
-    .cross::before, .cross::after {
-      content: "";
-      position: absolute;
-      left: 5px;
-      top: 0;
-      width: 2px;
-      height: 12px;
-      background: var(--orphan);
-      transform-origin: center;
-    }
-    .cross::before { transform: rotate(45deg); }
-    .cross::after { transform: rotate(-45deg); }
-    .tooltip {
-      position: fixed;
-      z-index: 10;
-      max-width: 460px;
-      padding: 10px 12px;
-      background: rgba(31, 27, 22, 0.94);
-      color: #f8f4ec;
-      border-radius: 10px;
-      pointer-events: none;
-      opacity: 0;
-      transform: translateY(4px);
-      transition: opacity 120ms ease, transform 120ms ease;
-      white-space: pre-wrap;
-      font-size: 12px;
-      line-height: 1.45;
-    }
-    .tooltip.visible {
-      opacity: 1;
-      transform: translateY(0);
-    }
-    .footer {
-      padding: 0 14px 14px;
-      color: var(--muted);
-      font-size: 12px;
-    }
-    body.compact-ui .controls-extra { display: none; }
-  </style>
-</head>
-<body>
-  <div class="page">
-    <div class="panel controls">
-      <div class="controls-header">
-        <div class="row">
-          <strong id="title">Trace Timeline</strong>
-          <span id="summary" style="color: var(--muted); font-size: 13px;"></span>
-        </div>
-        <button id="toggleControls" class="btn-subtle btn-icon" type="button">▴</button>
-      </div>
-      <div class="controls-extra" id="controlsExtra">
-        <div class="row">
-          <label>Filter <input id="filterText" type="text" placeholder="sample id / source / status / trace"></label>
-          <label>Attempt
-            <select id="attemptValue">
-              <option value="all">all attempt segments</option>
-            </select>
-          </label>
-          <label>Attempt View
-            <select id="attemptMode">
-              <option value="filter">filter segments</option>
-              <option value="highlight">highlight segments</option>
-            </select>
-          </label>
-          <label>View
-            <select id="viewMode">
-              <option value="collapsed">collapsed</option>
-              <option value="expanded">expanded</option>
-            </select>
-          </label>
-          <label>Sort
-            <select id="sortMode">
-              <option value="start">start</option>
-              <option value="duration">duration</option>
-              <option value="reward">reward</option>
-              <option value="lanes">lane count</option>
-              <option value="open">ongoing event count</option>
-            </select>
-          </label>
-          <label><input id="sortDesc" type="checkbox"> desc</label>
-          <label><input id="showEvents" type="checkbox" checked> instant events</label>
-          <label><input id="showOpenSpans" type="checkbox" checked> ongoing events</label>
-          <label><input id="showOrphans" type="checkbox" checked> unmatched ends</label>
-          <label>Lane Height <input id="laneHeight" type="range" min="14" max="28" value="18"></label>
-          <button id="fitAll" class="btn-subtle" type="button">fit all</button>
-          <button id="fitRunning" class="btn-subtle" type="button">fit filtered</button>
-        </div>
-        <div class="row">
-          <span style="color: var(--muted); font-size: 12px;">
-            drag = pan, wheel = zoom, click = set cursor and select item
-          </span>
-          <span id="cursorText" style="color: var(--accent); font-size: 12px;"></span>
-        </div>
-      </div>
-    </div>
-
-    <div class="panel stats">
-      <div id="stats" class="stats-grid"></div>
-    </div>
-
-    <div class="panel timeline-panel">
-      <div id="legend" class="legend"></div>
-      <div id="viewport" class="viewport">
-        <div id="canvasWrap" class="canvas-wrap">
-          <canvas id="timelineCanvas"></canvas>
-        </div>
-      </div>
-      <div class="footer" id="footer"></div>
-    </div>
-  </div>
-  <div id="tooltip" class="tooltip"></div>
-
-  <script>
-    const CACHE_FILE = "__CACHE_FILE__";
-    const LABEL_WIDTH = 320;
-    const AXIS_HEIGHT = 34;
-    const ROW_PADDING = 4;
-    const state = {
-      rawRows: [],
-      rows: [],
-      rowLayouts: [],
-      globalStart: 0,
-      globalEnd: 1,
-      viewStart: 0,
-      viewEnd: 1,
-      laneHeight: 18,
-      viewMode: 'collapsed',
-      cursorTime: null,
-      hoveredItem: null,
-      selectedItem: null,
-      dragging: false,
-      dragStartX: 0,
-      dragViewStart: 0,
-      dragViewEnd: 1,
-      rafScheduled: false,
-      compactUI: false,
-      expandedRows: new Set(),
-    };
-
-    const viewport = document.getElementById('viewport');
-    const canvasWrap = document.getElementById('canvasWrap');
-    const canvas = document.getElementById('timelineCanvas');
-    const tooltip = document.getElementById('tooltip');
-    const ctx = canvas.getContext('2d');
-
-    function niceDuration(seconds) {
-      if (!isFinite(seconds)) return 'n/a';
-      if (seconds < 1) return `${(seconds * 1000).toFixed(1)} ms`;
-      if (seconds < 60) return `${seconds.toFixed(2)} s`;
-      const minutes = Math.floor(seconds / 60);
-      const remain = seconds - minutes * 60;
-      return `${minutes}m ${remain.toFixed(1)}s`;
-    }
-
-    function niceNumber(value) {
-      if (value == null || Number.isNaN(value)) return 'n/a';
-      if (typeof value === 'number') return value.toLocaleString();
-      return `${value}`;
-    }
-
-    function escapeHtml(value) {
-      return `${value ?? ''}`
-        .replaceAll('&', '&amp;')
-        .replaceAll('<', '&lt;')
-        .replaceAll('>', '&gt;')
-        .replaceAll('"', '&quot;')
-        .replaceAll("'", '&#39;');
-    }
-
-    function hashHue(name) {
-      let hash = 0;
-      for (let idx = 0; idx < name.length; idx += 1) {
-        hash = ((hash << 5) - hash + name.charCodeAt(idx)) | 0;
-      }
-      return ((hash % 360) + 360) % 360;
-    }
-
-    function hashColor(name, alpha = 1) {
-      const hue = hashHue(name);
-      return `hsla(${hue} 55% 46% / ${alpha})`;
-    }
-
-    function pointColor(name, alpha = 1) {
-      const hue = hashHue(name);
-      return `hsla(${hue} 70% 24% / ${alpha})`;
-    }
-
-    function pointShape(name) {
-      const value = `${name ?? ''}`.toLowerCase();
-      if (value.includes('schedule') || value.includes('enqueue') || value.includes('dequeue') || value.includes('buffer')) {
-        return 'triangle';
-      }
-      if (value.includes('assign') || value.includes('complete') || value.includes('selected') || value.includes('finish')) {
-        return 'diamond';
-      }
-      if (value.includes('create') || value.includes('attempt') || value.includes('start')) {
-        return 'square';
-      }
-      return 'circle';
-    }
-
-    function itemStart(item) {
-      return item.start_ts ?? item.ts ?? 0;
-    }
-
-    function itemDisplayEnd(item) {
-      return item.display_end_ts ?? item.end_ts ?? item.ts ?? item.start_ts ?? 0;
-    }
-
-    function collapsedTextSafeEnd(row, item) {
-      let safeEnd = itemDisplayEnd(item);
-      if (!item.span_id) return safeEnd;
-      for (const candidate of row.items) {
-        if (candidate.type !== 'span') continue;
-        if (candidate.parent_span_id !== item.span_id) continue;
-        const childStart = itemStart(candidate);
-        if (childStart > item.start_ts && childStart < safeEnd) {
-          safeEnd = childStart;
-        }
-      }
-      return safeEnd;
-    }
-
-    function rowLabelParts(row) {
-      const sampleId = row.sample_index ?? row.row_id;
-      const expandMark = rowExpanded(row) ? '▾' : '▸';
-      return `${expandMark} #${sampleId} | ${row.source ?? 'unknown'}`;
-    }
-
-    function rowSummaryItem(row) {
-      const attempts = [...new Set((row.items || []).map(item => item.attempt).filter(v => v != null))].sort((a, b) => a - b);
-      return {
-        type: 'row_summary',
-        name: 'sample summary',
-        attempt: attempts.length > 1 ? attempts.join(',') : (attempts[0] ?? 'n/a'),
-        attrs: {
-          duration: niceDuration(row.duration ?? 0),
-          reward: typeof row.reward === 'number' ? row.reward.toFixed(3) : niceNumber(row.reward),
-          attempts: attempts.join(','),
-          status: row.status ?? 'unknown',
-          label: row.label ?? 'n/a',
-          trace_id: row.trace_id ?? 'n/a',
-        },
-      };
-    }
-
-    function sampleMatches(row, filterText) {
-      if (!filterText) return true;
-      const haystack = [
-        row.sample_index,
-        row.row_id,
-        row.source,
-        row.status,
-        row.label,
-        row.trace_id,
-      ].map(v => `${v ?? ''}`.toLowerCase()).join(' ');
-      return haystack.includes(filterText.toLowerCase());
-    }
-
-    function currentAttemptValue() {
-      return document.getElementById('attemptValue').value;
-    }
-
-    function currentAttemptMode() {
-      return document.getElementById('attemptMode').value;
-    }
-
-    function attemptMatches(item) {
-      const attemptValue = currentAttemptValue();
-      if (attemptValue === 'all') return true;
-      return Number(item.attempt) === Number(attemptValue);
-    }
-
-    function rowExpanded(row) {
-      return state.viewMode === 'expanded' || state.expandedRows.has(row.row_id);
-    }
-
-    function rowIsActive(row) {
-      return state.selectedItem?.row === row || state.hoveredItem?.row === row;
-    }
-
-    function rowAnalysisItems(row) {
-      const items = visibleItems(row);
-      const attemptValue = currentAttemptValue();
-      if (attemptValue === 'all') return items;
-      return items.filter(item => attemptMatches(item));
-    }
-
-    function visibleItems(row) {
-      const showEvents = document.getElementById('showEvents').checked;
-      const showOpenSpans = document.getElementById('showOpenSpans').checked;
-      const showOrphans = document.getElementById('showOrphans').checked;
-      return row.items.filter(item => {
-        const attrs = item.type === 'span' ? spanFlatAttrs(item) : (item.attrs || {});
-        if (attrs.timeline_pd_virtual_role && !rowExpanded(row)) return false;
-        if (item.state === 'point_event') return showEvents;
-        if (item.state === 'open_span') return showOpenSpans;
-        if (item.state === 'orphan_end') return showOrphans;
-        return true;
-      });
-    }
-
-    function rowHasMatchingAttempt(row) {
-      const attemptValue = currentAttemptValue();
-      if (attemptValue === 'all') return true;
-      return row.items.some(item => attemptMatches(item));
-    }
-
-    function rowVisibleLaneCount(row) {
-      const items = rowAnalysisItems(row);
-      return 1 + Math.max(0, ...items.map(item => item.lane || 0));
-    }
-
-    function itemVisualAlpha(item) {
-      const attemptValue = currentAttemptValue();
-      const attemptMode = currentAttemptMode();
-      if (attemptValue === 'all' || attemptMode !== 'highlight') return 1;
-      return attemptMatches(item) ? 1 : 0.18;
-    }
-
-    function itemIsFocusedAttempt(item) {
-      const attemptValue = currentAttemptValue();
-      return attemptValue !== 'all' && attemptMatches(item);
-    }
-
-    function drawnItemsForRow(row) {
-      return currentAttemptMode() === 'filter' ? rowAnalysisItems(row) : visibleItems(row);
-    }
-
-    function buildPointClusters(row, layout, scrollTop, items) {
-      const laneGroups = new Map();
-      const pointItems = items.filter(item => item.type !== 'span');
-      for (const item of pointItems) {
-        const lane = rowExpanded(row) ? (item.lane || 0) : 0;
-        if (!laneGroups.has(lane)) laneGroups.set(lane, []);
-        laneGroups.get(lane).push(item);
-      }
-
-      const clusters = [];
-      const threshold = 10;
-      for (const [lane, laneItems] of laneGroups.entries()) {
-        laneItems.sort((a, b) => itemStart(a) - itemStart(b));
-        let current = null;
-        for (const item of laneItems) {
-          const { y, h } = itemGeometry(layout, item, scrollTop);
-          const x = timeToX(item.ts ?? item.start_ts ?? 0);
-          if (!current || (x - current.maxX) > threshold) {
-            current = { lane, items: [item], xSum: x, minX: x, maxX: x, y, h };
-            clusters.push(current);
-            continue;
-          }
-          current.items.push(item);
-          current.xSum += x;
-          current.maxX = x;
-        }
-      }
-
-      return clusters.map(cluster => {
-        const x = cluster.xSum / Math.max(1, cluster.items.length);
-        if (cluster.items.length === 1) {
-          return { ...cluster, x, item: cluster.items[0], kind: 'single_point' };
-        }
-        const counts = new Map();
-        cluster.items.forEach(item => counts.set(item.name, (counts.get(item.name) || 0) + 1));
-        const countEntries = Array.from(counts.entries()).sort((a, b) => {
-          if (b[1] !== a[1]) return b[1] - a[1];
-          return a[0].localeCompare(b[0]);
-        });
-        return {
-          ...cluster,
-          x,
-          kind: 'point_cluster',
-          item: {
-            type: 'point_cluster',
-            state: 'point_cluster',
-            name: 'point cluster',
-            attempt: 'mixed',
-            attrs: { total: cluster.items.length, names: Object.fromEntries(countEntries) },
-            ts: cluster.items[0].ts ?? cluster.items[0].start_ts ?? 0,
-          },
-        };
-      });
-    }
-
-    function legendSpanNames() {
-      const counts = new Map();
-      for (const row of state.rows) {
-        for (const item of rowAnalysisItems(row)) {
-          if (item.type !== 'span') continue;
-          if (spanFlatAttrs(item).timeline_pd_virtual_role) continue;
-          const name = `${item.name ?? 'event'}`;
-          const current = counts.get(name) ?? { name, count: 0 };
-          current.count += 1;
-          counts.set(name, current);
-        }
-      }
-      return [...counts.values()]
-        .sort((a, b) => (b.count - a.count) || a.name.localeCompare(b.name))
-        .slice(0, 16);
-    }
-
-    function updateLegend() {
-      const legend = document.getElementById('legend');
-      if (!legend) return;
-      const names = legendSpanNames();
-      if (!names.length) {
-        legend.innerHTML = '<span class="chip">no span events</span>';
-        return;
-      }
-      // Check whether any row has PD data
-      let hasPD = false;
-      outer: for (const row of state.rows) {
-        for (const item of rowAnalysisItems(row)) {
-          if (item.type === 'span' && pdPhases(spanFlatAttrs(item))) { hasPD = true; break outer; }
-        }
-      }
-      let html = names.map((entry) => (
-        `<span class="chip" title="count=${entry.count}">` +
-        `<span class="bar" style="background: ${hashColor(entry.name, 0.9)}"></span>` +
-        `${escapeHtml(entry.name)}` +
-        `</span>`
-      )).join('');
-      if (hasPD) {
-        const pdLegend = [
-          ['prefill_forward', '[P] prefill fwd'],
-          ['prefill_bootstrap_queue', '[P] bootstrap queue'],
-          ['prefill_transfer_queue', '[P] transfer'],
-          ['bootstrap', '[P] bootstrap'],
-          ['alloc_waiting', '[P] alloc wait'],
-          ['decode_forward', '[D] decode fwd'],
-          ['decode_transfer', '[D] kv transfer'],
-          ['decode_prealloc', '[D] prealloc'],
-        ];
-        html += '<span style="color:var(--muted);margin-left:8px;font-size:11px">PD (collapsed: P over D; expanded: separate P/D lanes):</span>';
-        for (const [phase, label] of pdLegend) {
-          html += `<span class="chip" style="font-size:11px">` +
-            `<span class="bar" style="background:${pdPhaseColor(phase, 0.85)}"></span>` +
-            `${label}</span>`;
-        }
-      }
-      legend.innerHTML = html;
-    }
-
-    function applyFilterAndSort() {
-      const filterText = document.getElementById('filterText').value.trim();
-      const sortMode = document.getElementById('sortMode').value;
-      const sortDesc = document.getElementById('sortDesc').checked;
-      state.rows = state.rawRows.filter(row => sampleMatches(row, filterText));
-      if (currentAttemptMode() === 'filter') {
-        state.rows = state.rows.filter(row => rowHasMatchingAttempt(row));
-      }
-      state.rows.sort((a, b) => {
-        let delta = (a.start ?? 0) - (b.start ?? 0);
-        if (sortMode === 'duration') delta = (a.duration ?? 0) - (b.duration ?? 0);
-        if (sortMode === 'reward') delta = (Number(a.reward) || 0) - (Number(b.reward) || 0);
-        if (sortMode === 'lanes') delta = (a.lane_count ?? 0) - (b.lane_count ?? 0);
-        if (sortMode === 'open') delta = (a.open_span_count ?? 0) - (b.open_span_count ?? 0);
-        return sortDesc ? -delta : delta;
-      });
-      layout();
-      updateLegend();
-      updateStats();
-      scheduleDraw();
-    }
-
-    function computeRowLayouts() {
-      const layouts = [];
-      let cursor = AXIS_HEIGHT;
-      for (const row of state.rows) {
-        const collapsedHeight = state.laneHeight + ROW_PADDING * 2;
-        const expandedHeight = rowVisibleLaneCount(row) * state.laneHeight + ROW_PADDING * 2;
-        const rowHeight = Math.max(
-          state.laneHeight + ROW_PADDING * 2,
-          rowExpanded(row) ? expandedHeight : collapsedHeight,
-        );
-        layouts.push({ top: cursor, height: rowHeight });
-        cursor += rowHeight;
-      }
-      state.rowLayouts = layouts;
-      return cursor;
-    }
-
-    function layout() {
-      const totalHeight = computeRowLayouts();
-      canvasWrap.style.height = `${Math.max(totalHeight + 8, 120)}px`;
-      const width = Math.max(720, viewport.clientWidth - 2);
-      const height = Math.max(240, viewport.clientHeight - 2);
-      canvas.width = width * window.devicePixelRatio;
-      canvas.height = height * window.devicePixelRatio;
-      canvas.style.width = `${width}px`;
-      canvas.style.height = `${height}px`;
-      ctx.setTransform(window.devicePixelRatio, 0, 0, window.devicePixelRatio, 0, 0);
-      updateSummary();
-    }
-
-    function updateSummary() {
-      document.getElementById('summary').textContent =
-        `${state.rows.length} filtered / ${state.rawRows.length} total`;
-      document.getElementById('footer').textContent =
-        `view range: ${niceDuration(state.viewEnd - state.viewStart)} | global range: ${niceDuration(state.globalEnd - state.globalStart)}`;
-    }
-
-    function timeToX(timestamp) {
-      const width = canvas.clientWidth - LABEL_WIDTH;
-      return LABEL_WIDTH + ((timestamp - state.viewStart) / (state.viewEnd - state.viewStart)) * width;
-    }
-
-    function xToTime(x) {
-      const width = canvas.clientWidth - LABEL_WIDTH;
-      const ratio = (x - LABEL_WIDTH) / width;
-      return state.viewStart + ratio * (state.viewEnd - state.viewStart);
-    }
-
-    function laneTop(layout, lane, scrollTop) {
-      return layout.top - scrollTop + ROW_PADDING + lane * state.laneHeight;
-    }
-
-    function itemGeometry(layout, item, scrollTop) {
-      const rowIdx = state.rowLayouts.indexOf(layout);
-      const row = rowIdx >= 0 ? state.rows[rowIdx] : null;
-      if (row && rowExpanded(row)) {
-        return {
-          y: laneTop(layout, item.lane || 0, scrollTop) + 2,
-          h: Math.max(10, state.laneHeight - 4),
-        };
-      }
-
-      const baseY = layout.top - scrollTop + ROW_PADDING + 2;
-      return {
-        y: baseY,
-        h: Math.max(10, layout.height - ROW_PADDING * 2 - 4),
-      };
-    }
-
-    function drawAxis() {
-      const width = canvas.clientWidth;
-      const height = canvas.height / window.devicePixelRatio;
-      ctx.fillStyle = '#fffaf2';
-      ctx.fillRect(0, 0, width, AXIS_HEIGHT);
-      ctx.fillStyle = '#6f675d';
-      ctx.font = '12px IBM Plex Sans, sans-serif';
-      ctx.textBaseline = 'top';
-      ctx.fillText('sample', 12, 10);
-      const range = state.viewEnd - state.viewStart;
-      const targetTicks = 8;
-      const rawStep = range / targetTicks;
-      const power = Math.pow(10, Math.floor(Math.log10(Math.max(rawStep, 1e-6))));
-      const candidates = [1, 2, 5, 10];
-      let step = power;
-      for (const candidate of candidates) {
-        const maybe = candidate * power;
-        if (maybe >= rawStep) {
-          step = maybe;
-          break;
-        }
-      }
-      const firstTick = Math.ceil(state.viewStart / step) * step;
-      ctx.strokeStyle = '#d8cebf';
-      ctx.fillStyle = '#6f675d';
-      for (let tick = firstTick; tick <= state.viewEnd; tick += step) {
-        const x = timeToX(tick);
-        ctx.beginPath();
-        ctx.moveTo(x, AXIS_HEIGHT - 12);
-        ctx.lineTo(x, height);
-        ctx.stroke();
-        ctx.fillText(niceDuration(tick - state.globalStart), x + 4, 10);
-      }
-      ctx.beginPath();
-      ctx.moveTo(0, AXIS_HEIGHT - 0.5);
-      ctx.lineTo(width, AXIS_HEIGHT - 0.5);
-      ctx.strokeStyle = '#bdb3a3';
-      ctx.stroke();
-    }
-
-    function lowerBoundRows(targetY) {
-      let left = 0;
-      let right = state.rowLayouts.length;
-      while (left < right) {
-        const mid = (left + right) >> 1;
-        const layout = state.rowLayouts[mid];
-        if (layout.top + layout.height < targetY) {
-          left = mid + 1;
-        } else {
-          right = mid;
-        }
-      }
-      return left;
-    }
-
-    function visibleRowRange() {
-      const width = canvas.clientWidth;
-      const visibleStartY = viewport.scrollTop;
-      const visibleEndY = visibleStartY + viewport.clientHeight;
-      const startRow = Math.max(0, lowerBoundRows(visibleStartY) - 1);
-      let endRow = startRow;
-      while (endRow < state.rowLayouts.length && state.rowLayouts[endRow].top <= visibleEndY + state.laneHeight * 2) {
-        endRow += 1;
-      }
-      return { width, startRow, endRow: Math.max(startRow - 1, endRow - 1) };
-    }
-
-    // PD phase color palette (warm earth tones matching the viewer aesthetic)
-    const PD_COLORS = {
-      prefill_bootstrap_queue: 'rgba(180, 160, 120, A)',  // muted sand
-      bootstrap:               'rgba(140, 120, 90, A)',   // dark sand
-      alloc_waiting:           'rgba(160, 140, 110, A)',  // warm grey
-      prefill_forward:         'rgba(70, 140, 180, A)',   // steel blue (prefill)
-      prefill_transfer_queue:  'rgba(200, 160, 60, A)',   // amber (transfer)
-      decode_prealloc:         'rgba(160, 140, 110, A)',  // warm grey
-      decode_transfer:         'rgba(200, 160, 60, A)',   // amber (transfer)
-      decode_forward:          'rgba(80, 170, 100, A)',   // green (decode)
-    };
-    function pdPhaseColor(phase, alpha) {
-      return (PD_COLORS[phase] || 'rgba(128,128,128,A)').replace(/A/g, alpha.toFixed(2));
-    }
-
-    // For span items, attrs is {start_attrs, end_attrs}. Flatten into a single dict.
-    function spanFlatAttrs(item) {
-      const raw = item.attrs || {};
-      if (raw.start_attrs || raw.end_attrs) {
-        return Object.assign({}, raw.start_attrs || {}, raw.end_attrs || {});
-      }
-      return raw;
-    }
-
-    // Build ordered PD phase segments from span attrs.
-    // Returns [{phase, duration}] or null if no PD data.
-    function pdPhases(attrs) {
-      if (!attrs) return null;
-      const segs = [];
-      const push = (phase, key) => {
-        const v = attrs[`pd_${key}`];
-        if (v != null && v > 0) segs.push({ phase, duration: v });
-      };
-      // P-side phases (sequential order)
-      push('prefill_bootstrap_queue', 'prefill_bootstrap_queue_duration');
-      push('bootstrap',               'bootstrap_duration');
-      push('alloc_waiting',           'alloc_waiting_duration');
-      push('prefill_forward',         'prefill_forward_duration');
-      push('prefill_transfer_queue',  'prefill_transfer_queue_duration');
-      // D-side phases (sequential order)
-      push('decode_prealloc',  'decode_prealloc_duration');
-      push('decode_transfer',  'decode_transfer_duration');
-      push('decode_forward',   'decode_forward_duration');
-      return segs.length > 0 ? segs : null;
-    }
-
-    function splitPDPhases(phases) {
-      if (!phases) {
-        return {
-          pPhases: [],
-          dPhases: [],
-          pTotal: 0,
-          dTotal: 0,
-          totalDuration: 0,
-        };
-      }
-      const pPhases = phases.filter((p) =>
-        p.phase.startsWith('prefill_') || p.phase === 'bootstrap' || p.phase === 'alloc_waiting');
-      const dPhases = phases.filter((p) => p.phase.startsWith('decode_'));
-      const pTotal = pPhases.reduce((sum, p) => sum + p.duration, 0);
-      const dTotal = dPhases.reduce((sum, p) => sum + p.duration, 0);
-      return {
-        pPhases,
-        dPhases,
-        pTotal,
-        dTotal,
-        totalDuration: pTotal + dTotal,
-      };
-    }
-
-    function pdRenderInfo(item, expanded) {
-      const attrs = spanFlatAttrs(item);
-      const pdVirtualRole = attrs.timeline_pd_virtual_role || null;
-      const phases = pdPhases(attrs);
-      if (!phases || (expanded && !pdVirtualRole)) return null;
-      const {
-        pPhases,
-        dPhases,
-        pTotal,
-        dTotal,
-        totalDuration,
-      } = splitPDPhases(phases);
-      const itemDuration = Math.max(0, itemDisplayEnd(item) - itemStart(item));
-      const scaleDuration = Math.max(totalDuration, itemDuration, 1e-9);
-      let filledDuration = totalDuration;
-      if (expanded && pdVirtualRole === 'prefill') filledDuration = pTotal;
-      if (expanded && pdVirtualRole === 'decode') filledDuration = dTotal;
-      return {
-        attrs,
-        pdVirtualRole,
-        phases,
-        pPhases,
-        dPhases,
-        pTotal,
-        dTotal,
-        totalDuration,
-        scaleDuration,
-        filledDuration,
-      };
-    }
-
-    function drawPDPhaseRow(segs, scaleDuration, x1, spanWidth, ry, rh, alpha) {
-      if (!segs.length || scaleDuration <= 0) return;
-      let cx = x1;
-      for (const seg of segs) {
-        const w = Math.max(1, (seg.duration / scaleDuration) * spanWidth);
-        ctx.fillStyle = pdPhaseColor(seg.phase, 0.85 * alpha);
-        ctx.fillRect(cx, ry, w, rh);
-        if (cx > x1) {
-          ctx.fillStyle = 'rgba(0,0,0,0.25)';
-          ctx.fillRect(cx, ry, 0.5, rh);
-        }
-        cx += w;
-      }
-    }
-
-    function drawSpan(item, x1, x2, y, h, expanded) {
-      const attrs = spanFlatAttrs(item);
-      const colorKey = attrs.timeline_pd_parent_name || item.name;
-      const alpha = itemVisualAlpha(item);
-      const color = hashColor(colorKey, item.state === 'open_span' ? 0.25 * alpha : 0.9 * alpha);
-      const border = hashColor(colorKey, Math.max(0.22, alpha));
-      const spanWidth = Math.max(1.5, x2 - x1);
-      const pdVirtualRole = attrs.timeline_pd_virtual_role || null;
-
-      if (item.state === 'open_span') {
-        ctx.save();
-        ctx.fillStyle = color;
-        ctx.fillRect(x1, y, Math.max(2, x2 - x1), h);
-        ctx.strokeStyle = border;
-        ctx.setLineDash([6, 4]);
-        ctx.strokeRect(x1 + 0.5, y + 0.5, Math.max(1, x2 - x1 - 1), Math.max(1, h - 1));
-        ctx.setLineDash([]);
-        ctx.restore();
-        const fadeWidth = Math.min(18, Math.max(6, x2 - x1));
-        const gradient = ctx.createLinearGradient(x2 - fadeWidth, 0, x2, 0);
-        gradient.addColorStop(0, hashColor(colorKey, 0));
-        gradient.addColorStop(1, hashColor(colorKey, Math.max(0.22, alpha)));
-        ctx.fillStyle = gradient;
-        ctx.fillRect(Math.max(x1, x2 - fadeWidth), y, fadeWidth, h);
-        if (itemIsFocusedAttempt(item)) {
-          ctx.save();
-          ctx.strokeStyle = '#1f4b3a';
-          ctx.lineWidth = 1.5;
-          ctx.strokeRect(x1 + 0.5, y + 0.5, Math.max(1, x2 - x1 - 1), Math.max(1, h - 1));
-          ctx.restore();
-        }
-        return;
-      }
-
-      // Draw PD phase sub-bars when disaggregation data is available
-      const pdInfo = pdRenderInfo(item, expanded);
-      if (pdInfo && spanWidth > 6) {
-        const {
-          pPhases,
-          dPhases,
-          totalDuration,
-          scaleDuration,
-        } = pdInfo;
-        const hasP = pPhases.length > 0;
-        const hasD = dPhases.length > 0;
-
-        // Base bar at reduced opacity
-        ctx.fillStyle = hashColor(item.name, 0.15 * alpha);
-        ctx.fillRect(x1, y, spanWidth, h);
-
-        if (!expanded) {
-          // Collapsed: keep a single row, with P painted over D so short prefill
-          // timings remain visible instead of being hidden by longer decode phases.
-          const subY = y + 2;
-          const subH = Math.max(4, h - 4);
-          if (hasD) drawPDPhaseRow(dPhases, scaleDuration, x1, spanWidth, subY, subH, 0.62 * alpha);
-          if (hasP) drawPDPhaseRow(pPhases, scaleDuration, x1, spanWidth, subY, subH, alpha);
-        } else {
-          const subY = y + 2;
-          const subH = Math.max(4, h - 4);
-          if (pdVirtualRole === 'prefill' && hasP) {
-            drawPDPhaseRow(pPhases, scaleDuration, x1, spanWidth, subY, subH, alpha);
-          } else if (pdVirtualRole === 'decode' && hasD) {
-            drawPDPhaseRow(dPhases, scaleDuration, x1, spanWidth, subY, subH, 0.9 * alpha);
-          }
-        }
-
-        // Border around the full span
-        ctx.strokeStyle = border;
-        ctx.lineWidth = 0.5;
-        ctx.strokeRect(x1 + 0.5, y + 0.5, Math.max(1, spanWidth - 1), Math.max(1, h - 1));
-      } else {
-        ctx.fillStyle = color;
-        ctx.fillRect(x1, y, spanWidth, h);
-      }
-
-      if (itemIsFocusedAttempt(item)) {
-        ctx.save();
-        ctx.strokeStyle = '#1f4b3a';
-        ctx.lineWidth = 1.5;
-        ctx.strokeRect(x1 + 0.5, y + 0.5, Math.max(1, x2 - x1 - 1), Math.max(1, h - 1));
-        ctx.restore();
-      }
-    }
-
-    function drawPoint(item, x, y, h) {
-      const midY = y + h / 2;
-      if (item.state === 'orphan_end') {
-        ctx.strokeStyle = getComputedStyle(document.documentElement).getPropertyValue('--orphan').trim();
-        ctx.globalAlpha = itemVisualAlpha(item);
-        ctx.lineWidth = 2;
-        ctx.beginPath();
-        ctx.moveTo(x - 4, midY - 4);
-        ctx.lineTo(x + 4, midY + 4);
-        ctx.moveTo(x + 4, midY - 4);
-        ctx.lineTo(x - 4, midY + 4);
-        ctx.stroke();
-        ctx.globalAlpha = 1;
-        return;
-      }
-      const color = pointColor(item.name, itemVisualAlpha(item));
-      const shape = pointShape(item.name);
-      ctx.strokeStyle = color;
-      ctx.lineWidth = 1.5;
-      ctx.beginPath();
-      ctx.moveTo(x, y);
-      ctx.lineTo(x, y + h);
-      ctx.stroke();
-      ctx.fillStyle = color;
-      if (shape === 'triangle') {
-        ctx.beginPath();
-        ctx.moveTo(x, midY - 5);
-        ctx.lineTo(x + 5, midY + 4);
-        ctx.lineTo(x - 5, midY + 4);
-        ctx.closePath();
-        ctx.fill();
-        return;
-      }
-      if (shape === 'square') {
-        ctx.fillRect(x - 4, midY - 4, 8, 8);
-        return;
-      }
-      if (shape === 'circle') {
-        ctx.beginPath();
-        ctx.arc(x, midY, 4.5, 0, Math.PI * 2);
-        ctx.fill();
-        return;
-      }
-      ctx.beginPath();
-      ctx.moveTo(x, midY - 5);
-      ctx.lineTo(x + 5, midY);
-      ctx.lineTo(x, midY + 5);
-      ctx.lineTo(x - 5, midY);
-      ctx.closePath();
-      ctx.fill();
-    }
-
-    function drawPointCluster(cluster) {
-      if (cluster.kind === 'single_point') {
-        drawPoint(cluster.item, cluster.x, cluster.y, cluster.h);
-        return;
-      }
-      const midY = cluster.y + cluster.h / 2;
-      ctx.fillStyle = 'rgba(32, 28, 24, 0.88)';
-      ctx.strokeStyle = 'rgba(255,255,255,0.9)';
-      ctx.lineWidth = 1.2;
-      ctx.beginPath();
-      ctx.arc(cluster.x, midY, 7, 0, Math.PI * 2);
-      ctx.fill();
-      ctx.stroke();
-      ctx.fillStyle = '#f7f2ea';
-      ctx.font = '10px IBM Plex Sans, sans-serif';
-      ctx.textAlign = 'center';
-      ctx.textBaseline = 'middle';
-      const label = cluster.items.length > 9 ? '9+' : `${cluster.items.length}`;
-      ctx.fillText(label, cluster.x, midY + 0.5);
-      ctx.textAlign = 'start';
-      ctx.textBaseline = 'alphabetic';
-    }
-
-    function drawRows() {
-      const { width, startRow, endRow } = visibleRowRange();
-      if (endRow < startRow) return;
-      const scrollTop = viewport.scrollTop;
-      const showText = (state.viewEnd - state.viewStart) <= 120;
-
-      for (let rowIdx = startRow; rowIdx <= endRow; rowIdx += 1) {
-        const row = state.rows[rowIdx];
-        const layout = state.rowLayouts[rowIdx];
-        if (!row || !layout) continue;
-        const top = layout.top - scrollTop;
-        const activeRow = rowIsActive(row);
-        const expandedRow = rowExpanded(row) && state.viewMode !== 'expanded';
-        ctx.fillStyle = activeRow
-          ? 'rgba(181, 80, 42, 0.12)'
-          : expandedRow
-            ? 'rgba(47, 109, 155, 0.08)'
-            : rowIdx % 2 === 0
-              ? 'rgba(216, 206, 191, 0.10)'
-              : 'rgba(216, 206, 191, 0.03)';
-        ctx.fillRect(0, top, width, layout.height);
-        if (expandedRow) {
-          ctx.fillStyle = 'rgba(47, 109, 155, 0.10)';
-          ctx.fillRect(0, top, LABEL_WIDTH, layout.height);
-        }
-        const label = rowLabelParts(row);
-        ctx.save();
-        ctx.beginPath();
-        ctx.rect(0, top, LABEL_WIDTH - 8, layout.height);
-        ctx.clip();
-        ctx.textBaseline = 'middle';
-        ctx.fillStyle = activeRow ? '#6c341f' : '#5f574d';
-        ctx.font = '11px IBM Plex Sans, sans-serif';
-        ctx.fillText(label, 12, top + layout.height / 2);
-        ctx.restore();
-
-        if (rowExpanded(row)) {
-          for (let lane = 1; lane < rowVisibleLaneCount(row); lane += 1) {
-            const y = laneTop(layout, lane, scrollTop) - 1;
-            ctx.strokeStyle = 'rgba(189, 179, 163, 0.35)';
-            ctx.beginPath();
-            ctx.moveTo(LABEL_WIDTH, y);
-            ctx.lineTo(width, y);
-            ctx.stroke();
-          }
-        }
-
-        const drawnItems = drawnItemsForRow(row);
-        const spanItems = drawnItems.filter(item => item.type === 'span');
-        const pointClusters = buildPointClusters(row, layout, scrollTop, drawnItems);
-        ctx.save();
-        ctx.beginPath();
-        ctx.rect(LABEL_WIDTH, top, Math.max(1, width - LABEL_WIDTH), layout.height);
-        ctx.clip();
-        for (const item of spanItems) {
-          const { y, h } = itemGeometry(layout, item, scrollTop);
-          const start = itemStart(item);
-          const end = itemDisplayEnd(item);
-          if (end < state.viewStart || start > state.viewEnd) continue;
-          if (item.type === 'span') {
-            const x1 = Math.max(LABEL_WIDTH, timeToX(start));
-            const x2 = Math.min(width, timeToX(end));
-            drawSpan(item, x1, x2, y, h, rowExpanded(row));
-            if (state.selectedItem?.item === item) {
-              ctx.save();
-              ctx.strokeStyle = '#b5502a';
-              ctx.lineWidth = 2;
-              ctx.strokeRect(x1 + 0.5, y + 0.5, Math.max(1, x2 - x1 - 1), Math.max(1, h - 1));
-              ctx.restore();
-            }
-            const expanded = rowExpanded(row);
-            const labelEnd = !expanded
-              ? Math.min(end, collapsedTextSafeEnd(row, item))
-              : end;
-            const labelX2 = Math.min(width, timeToX(labelEnd));
-            const labelWidthAvailable = labelX2 - x1;
-            if (showText && labelWidthAvailable > 40) {
-              ctx.font = '10px IBM Plex Sans, sans-serif';
-              const label = (() => {
-                const a = spanFlatAttrs(item);
-                const pdParts = [];
-                if (a.prompt_tokens != null) pdParts.push(`P:${a.prompt_tokens}`);
-                if (a.completion_tokens != null) pdParts.push(`D:${a.completion_tokens}`);
-                if (a.cached_tokens != null && a.cached_tokens > 0) pdParts.push(`cache:${a.cached_tokens}`);
-                if (a.pd_prefill_forward_duration != null) pdParts.push(`pf:${(a.pd_prefill_forward_duration*1000).toFixed(0)}ms`);
-                if (a.pd_decode_forward_duration != null) pdParts.push(`df:${(a.pd_decode_forward_duration*1000).toFixed(0)}ms`);
-                if (a.pd_transfer_speed_gb_s != null) pdParts.push(`${a.pd_transfer_speed_gb_s.toFixed(1)}GB/s`);
-                  const pdSuffix = pdParts.length > 0 ? ` | ${pdParts.join(' ')}` : '';
-                  return `${item.name} | attempt=${item.attempt}${pdSuffix}`;
-                })();
-              const labelWidth = ctx.measureText(label).width;
-              let clipX = x1;
-              let clipWidth = labelWidthAvailable;
-              let textX = x1 + 4;
-              const pdInfo = pdRenderInfo(item, expanded);
-              if (pdInfo) {
-                const spanWidth = Math.max(1.5, x2 - x1);
-                const filledWidth = Math.min(
-                  spanWidth,
-                  (pdInfo.filledDuration / pdInfo.scaleDuration) * spanWidth,
-                );
-                const safeX = Math.min(labelX2 - 8, x1 + filledWidth + 6);
-                const safeWidth = labelX2 - safeX;
-                if (labelWidth + 12 <= safeWidth) {
-                  clipX = safeX;
-                  clipWidth = safeWidth;
-                  textX = safeX + 4;
-                } else if (filledWidth > 12) {
-                  clipWidth = 0;
-                }
-              }
-              if (clipWidth > 0 && labelWidth + 12 <= clipWidth) {
-                ctx.fillStyle = itemVisualAlpha(item) < 1 ? 'rgba(255,255,255,0.45)' : 'rgba(255,255,255,0.86)';
-                ctx.font = '10px IBM Plex Sans, sans-serif';
-                ctx.save();
-                ctx.beginPath();
-                ctx.rect(clipX, y, clipWidth, h);
-                ctx.clip();
-                ctx.textBaseline = 'middle';
-                ctx.fillText(label, textX, y + h / 2 + 0.5);
-                ctx.restore();
-              }
-            }
-            continue;
-          }
-        }
-        pointClusters.forEach(cluster => drawPointCluster(cluster));
-        ctx.restore();
-      }
-    }
-
-    function drawCursor() {
-      if (state.cursorTime == null) return;
-      if (state.cursorTime < state.viewStart || state.cursorTime > state.viewEnd) return;
-      const x = timeToX(state.cursorTime);
-      ctx.beginPath();
-      ctx.moveTo(x, 0);
-      ctx.lineTo(x, canvas.height / window.devicePixelRatio);
-      ctx.lineWidth = 1;
-      ctx.strokeStyle = '#b5502a';
-      ctx.stroke();
-    }
-
-    function itemTooltipLines(row, item) {
-      const lines = [
-        `sample: ${row.sample_index ?? row.row_id}`,
-        `source: ${row.source ?? 'unknown'}`,
-        `name: ${item.name}`,
-        `attempt: ${item.attempt}`,
-      ];
-      if (item.type === 'row_summary') {
-        for (const [key, value] of Object.entries(item.attrs || {})) {
-          lines.push(`${key}: ${value}`);
-        }
-        return lines;
-      }
-      if (item.type === 'point_cluster') {
-        lines.push(`events: ${item.attrs?.total ?? 0}`);
-        Object.entries(item.attrs?.names || {}).slice(0, 8).forEach(([name, count]) => {
-          lines.push(`${name}: ${count}`);
-        });
-        return lines;
-      }
-      if (item.parent_span_name) lines.push(`parent: ${item.parent_span_name}`);
-      if (item.type === 'span') {
-        lines.push(`start: ${niceDuration(item.start_ts - state.globalStart)}`);
-        lines.push(`end: ${item.end_ts == null ? 'open' : niceDuration(item.end_ts - state.globalStart)}`);
-        if (item.end_ts == null) {
-          lines.push(`display_end: ${niceDuration(item.display_end_ts - state.globalStart)}`);
-        }
-        lines.push(`duration: ${item.end_ts == null ? 'open' : niceDuration(item.end_ts - item.start_ts)}`);
-      } else {
-        lines.push(`time: ${niceDuration((item.ts ?? 0) - state.globalStart)}`);
-      }
-      if (item.state === 'orphan_end') lines.push('warning: unmatched end event');
-      const attrs = spanFlatAttrs(item);
-      // Collect and group PD disaggregation attrs for structured display
-      const pdKeys = Object.keys(attrs).filter(k => k.startsWith('pd_'));
-      const otherKeys = Object.keys(attrs).filter(k => !k.startsWith('pd_') && !k.startsWith('timeline_'));
-      for (const key of otherKeys) {
-        const value = attrs[key];
-        lines.push(`${key}: ${typeof value === 'object' ? JSON.stringify(value) : value}`);
-      }
-      if (pdKeys.length > 0) {
-        const phases = pdPhases(attrs);
-        if (phases) {
-          const totalDur = phases.reduce((s, p) => s + p.duration, 0);
-          const pPhases = phases.filter(p => p.phase.startsWith('prefill_') || p.phase === 'bootstrap' || p.phase === 'alloc_waiting');
-          const dPhases = phases.filter(p => p.phase.startsWith('decode_'));
-          const fmtPhase = (p) => {
-            const ms = (p.duration * 1000).toFixed(1);
-            const pct = totalDur > 0 ? ((p.duration / totalDur) * 100).toFixed(0) : 0;
-            const nice = p.phase.replace(/_/g, ' ');
-            return `  ${nice}: ${ms}ms (${pct}%)`;
-          };
-          lines.push('── PD disaggregation ──');
-          if (pPhases.length > 0) {
-            lines.push(' [P] prefill instance');
-            pPhases.forEach(p => lines.push(fmtPhase(p)));
-          }
-          if (dPhases.length > 0) {
-            lines.push(' [D] decode instance');
-            dPhases.forEach(p => lines.push(fmtPhase(p)));
-          }
-          lines.push(`  total phases: ${(totalDur * 1000).toFixed(1)}ms`);
-        }
-        // Show non-duration PD fields (transfer speed, total MB, retry count)
-        for (const key of pdKeys) {
-          const value = attrs[key];
-          const label = key.replace(/^pd_/, '');
-          if (label.endsWith('_duration')) continue;  // already shown above
-          if (label.endsWith('_gb_s') && typeof value === 'number') {
-            lines.push(`  ${label}: ${value.toFixed(2)} GB/s`);
-          } else if (label.endsWith('_mb') && typeof value === 'number') {
-            lines.push(`  ${label}: ${value.toFixed(1)} MB`);
-          } else {
-            lines.push(`  ${label}: ${typeof value === 'object' ? JSON.stringify(value) : value}`);
-          }
-        }
-      }
-      return lines;
-    }
-
-    function updateTooltip() {
-      const activeItem = state.hoveredItem || state.selectedItem;
-      if (!activeItem) {
-        tooltip.classList.remove('visible');
-        return;
-      }
-      const { row, item, x, y } = activeItem;
-      tooltip.textContent = itemTooltipLines(row, item).join('\n');
-      tooltip.style.left = `${Math.min(window.innerWidth - 480, x + 14)}px`;
-      tooltip.style.top = `${Math.min(window.innerHeight - 200, y + 14)}px`;
-      tooltip.classList.add('visible');
-    }
-
-    function draw() {
-      ctx.clearRect(0, 0, canvas.clientWidth, canvas.height / window.devicePixelRatio);
-      drawRows();
-      drawAxis();
-      drawCursor();
-      updateTooltip();
-      updateSummary();
-    }
-
-    function scheduleDraw() {
-      if (state.rafScheduled) return;
-      state.rafScheduled = true;
-      requestAnimationFrame(() => {
-        state.rafScheduled = false;
-        draw();
-      });
-    }
-
-    function panTimelineByPixels(deltaX) {
-      const width = Math.max(1, canvas.clientWidth - LABEL_WIDTH);
-      const shift = (deltaX / width) * (state.viewEnd - state.viewStart);
-      let nextStart = state.viewStart + shift;
-      let nextEnd = state.viewEnd + shift;
-      if (nextStart < state.globalStart) {
-        nextEnd += state.globalStart - nextStart;
-        nextStart = state.globalStart;
-      }
-      if (nextEnd > state.globalEnd) {
-        nextStart -= nextEnd - state.globalEnd;
-        nextEnd = state.globalEnd;
-      }
-      state.viewStart = nextStart;
-      state.viewEnd = nextEnd;
-      scheduleDraw();
-    }
-
-    function setCompactUI(nextCompact) {
-      state.compactUI = nextCompact;
-      document.body.classList.toggle('compact-ui', nextCompact);
-      document.getElementById('toggleControls').textContent = nextCompact ? '▾' : '▴';
-      layout();
-      scheduleDraw();
-    }
-
-    function fitToRows(rows) {
-      if (!rows.length) return;
-      let minStart = rows[0].start;
-      let maxEnd = rows[0].end;
-      rows.forEach(row => {
-        minStart = Math.min(minStart, row.start);
-        maxEnd = Math.max(maxEnd, row.end);
-      });
-      const pad = Math.max((maxEnd - minStart) * 0.03, 0.001);
-      state.viewStart = minStart - pad;
-      state.viewEnd = maxEnd + pad;
-    }
-
-    function updateStats() {
-      const lifecycleStats = {
-        rows: state.rows.length,
-        notStarted: 0,
-        running: 0,
-        ended: 0,
-      };
-      const nameCounts = new Map();
-      const cursorTime = state.cursorTime ?? state.viewStart;
-      const pointTimeWindow = (state.viewEnd - state.viewStart) * 6 / Math.max(1, canvas.clientWidth - LABEL_WIDTH);
-
-      function addCount(name) {
-        nameCounts.set(name, (nameCounts.get(name) || 0) + 1);
-      }
-
-      state.rows.forEach(row => {
-        if (cursorTime < row.start) {
-          lifecycleStats.notStarted += 1;
-          return;
-        }
-        if (cursorTime > row.end) {
-          lifecycleStats.ended += 1;
-          return;
-        }
-        lifecycleStats.running += 1;
-
-        for (const item of rowAnalysisItems(row)) {
-          if (item.type === 'span') {
-            if (cursorTime >= item.start_ts && cursorTime <= itemDisplayEnd(item)) {
-              addCount(item.name);
-            }
-            continue;
-          }
-          const ts = item.ts ?? 0;
-          if (Math.abs(ts - cursorTime) <= pointTimeWindow) {
-            addCount(item.name);
-          }
-        }
-      });
-      document.getElementById('cursorText').textContent =
-        `cursor: ${niceDuration(cursorTime - state.globalStart)}`;
-      const lifecycleItems = [
-        { name: 'total samples', value: lifecycleStats.rows, kind: 'lifecycle' },
-        { name: 'not started', value: lifecycleStats.notStarted, kind: 'lifecycle' },
-        { name: 'running', value: lifecycleStats.running, kind: 'lifecycle' },
-        { name: 'ended', value: lifecycleStats.ended, kind: 'lifecycle' },
-      ];
-      const eventItems = Array.from(nameCounts.entries())
-        .sort((a, b) => {
-          if (b[1] !== a[1]) return b[1] - a[1];
-          return a[0].localeCompare(b[0]);
-        })
-        .map(([name, value]) => ({ name, value, kind: 'event' }));
-      const items = lifecycleItems.concat(eventItems);
-      document.getElementById('stats').innerHTML = items.map((item) => `
-        <div class="stat">
-          <div class="name">${
-            item.kind === 'event'
-              ? `<span class="swatch" style="background:${hashColor(item.name, 0.9)}"></span>${escapeHtml(item.name)}`
-              : escapeHtml(item.name)
-          }</div>
-          <div class="value">${item.value}</div>
-        </div>
-      `).join('');
-    }
-
-    function handleZoom(mouseX, deltaY) {
-      const focusTime = xToTime(mouseX);
-      const range = state.viewEnd - state.viewStart;
-      const scale = deltaY < 0 ? 0.9 : 1.1;
-      const newRange = Math.min(state.globalEnd - state.globalStart, Math.max(0.01, range * scale));
-      const ratio = (focusTime - state.viewStart) / range;
-      let nextStart = focusTime - ratio * newRange;
-      let nextEnd = nextStart + newRange;
-      if (nextStart < state.globalStart) {
-        nextEnd += state.globalStart - nextStart;
-        nextStart = state.globalStart;
-      }
-      if (nextEnd > state.globalEnd) {
-        nextStart -= nextEnd - state.globalEnd;
-        nextEnd = state.globalEnd;
-      }
-      state.viewStart = nextStart;
-      state.viewEnd = nextEnd;
-      scheduleDraw();
-    }
-
-    function rowIndexForCanvasY(canvasY) {
-      const absoluteY = canvasY + viewport.scrollTop;
-      const idx = lowerBoundRows(absoluteY);
-      if (idx >= state.rowLayouts.length) return -1;
-      const layout = state.rowLayouts[idx];
-      if (absoluteY >= layout.top && absoluteY <= layout.top + layout.height) return idx;
-      return -1;
-    }
-
-    function findItemAt(row, rowIdx, canvasX, canvasY) {
-      const layout = state.rowLayouts[rowIdx];
-      if (!layout) return null;
-      const localY = canvasY + viewport.scrollTop - layout.top;
-      const cursorTime = xToTime(canvasX);
-      const pointPixelTolerance = 8;
-      const hitItems = drawnItemsForRow(row);
-      const pointClusters = buildPointClusters(row, layout, 0, hitItems);
-      let closestPoint = null;
-      let closestPointDistance = Infinity;
-
-      for (const cluster of pointClusters) {
-        const pointY = (cluster.y - layout.top) + cluster.h / 2;
-        const dx = canvasX - cluster.x;
-        const dy = localY - pointY;
-        const radius = cluster.kind === 'point_cluster' ? 9 : pointPixelTolerance;
-        const distance = Math.hypot(dx, dy);
-        if (distance <= radius && distance < closestPointDistance) {
-          closestPoint = cluster.item;
-          closestPointDistance = distance;
-        }
-      }
-
-      if (closestPoint) return closestPoint;
-
-      for (const item of [...hitItems].reverse()) {
-        if (item.type !== 'span') continue;
-        const geometry = itemGeometry(layout, item, 0);
-        const itemTop = geometry.y - layout.top;
-        const itemBottom = itemTop + geometry.h;
-        if (localY < itemTop || localY > itemBottom) continue;
-        if (cursorTime >= item.start_ts && cursorTime <= itemDisplayEnd(item)) return item;
-      }
-
-      return null;
-    }
-
-    function attachEvents() {
-      document.getElementById('filterText').addEventListener('input', applyFilterAndSort);
-      document.getElementById('attemptValue').addEventListener('change', applyFilterAndSort);
-      document.getElementById('attemptMode').addEventListener('change', applyFilterAndSort);
-      document.getElementById('viewMode').addEventListener('change', (event) => {
-        state.viewMode = event.target.value;
-        layout();
-        updateLegend();
-        updateStats();
-        scheduleDraw();
-      });
-      document.getElementById('sortMode').addEventListener('change', applyFilterAndSort);
-      document.getElementById('sortDesc').addEventListener('change', applyFilterAndSort);
-      document.getElementById('showEvents').addEventListener('change', () => { updateLegend(); updateStats(); scheduleDraw(); });
-      document.getElementById('showOpenSpans').addEventListener('change', () => { layout(); updateLegend(); updateStats(); scheduleDraw(); });
-      document.getElementById('showOrphans').addEventListener('change', () => { updateLegend(); updateStats(); scheduleDraw(); });
-      document.getElementById('toggleControls').addEventListener('click', () => setCompactUI(!state.compactUI));
-      document.getElementById('laneHeight').addEventListener('input', (event) => {
-        state.laneHeight = Number(event.target.value);
-        layout();
-        scheduleDraw();
-      });
-      document.getElementById('fitAll').addEventListener('click', () => {
-        fitToRows(state.rawRows);
-        scheduleDraw();
-      });
-      document.getElementById('fitRunning').addEventListener('click', () => {
-        fitToRows(state.rows);
-        scheduleDraw();
-      });
-      viewport.addEventListener('scroll', () => scheduleDraw());
-      canvas.addEventListener('mousemove', (event) => {
-        if (state.dragging) {
-          const dx = event.clientX - state.dragStartX;
-          const dy = event.clientY - state.dragStartY;
-          if (!state.dragMoved && Math.hypot(dx, dy) > 3) {
-            state.dragMoved = true;
-          }
-          const pxToTime = (state.dragViewEnd - state.dragViewStart) / (canvas.clientWidth - LABEL_WIDTH);
-          const shift = dx * pxToTime;
-          let nextStart = state.dragViewStart - shift;
-          let nextEnd = state.dragViewEnd - shift;
-          if (nextStart < state.globalStart) {
-            nextEnd += state.globalStart - nextStart;
-            nextStart = state.globalStart;
-          }
-          if (nextEnd > state.globalEnd) {
-            nextStart -= nextEnd - state.globalEnd;
-            nextEnd = state.globalEnd;
-          }
-          state.viewStart = nextStart;
-          state.viewEnd = nextEnd;
-          scheduleDraw();
-          return;
-        }
-        const rect = canvas.getBoundingClientRect();
-        const canvasX = event.clientX - rect.left;
-        const canvasY = event.clientY - rect.top;
-        const rowIdx = rowIndexForCanvasY(canvasY);
-        if (canvasX < LABEL_WIDTH) {
-          let hoveredItem = null;
-          if (rowIdx >= 0 && rowIdx < state.rows.length) {
-            const row = state.rows[rowIdx];
-            hoveredItem = { row, item: rowSummaryItem(row), x: event.clientX, y: event.clientY };
-          }
-          const sameHovered = state.hoveredItem
-            && hoveredItem
-            && state.hoveredItem.row === hoveredItem.row
-            && state.hoveredItem.item.type === hoveredItem.item.type;
-          if (!sameHovered || (!hoveredItem && state.hoveredItem)) {
-            state.hoveredItem = hoveredItem;
-            updateTooltip();
-          } else if (hoveredItem) {
-            state.hoveredItem = hoveredItem;
-            updateTooltip();
-          }
-          return;
-        }
-        let hoveredItem = null;
-        if (rowIdx >= 0 && rowIdx < state.rows.length) {
-          const row = state.rows[rowIdx];
-          const item = findItemAt(row, rowIdx, canvasX, canvasY);
-          if (item && item.type !== 'span') hoveredItem = { row, item, x: event.clientX, y: event.clientY };
-        }
-        const sameHovered = state.hoveredItem
-          && hoveredItem
-          && state.hoveredItem.item === hoveredItem.item
-          && state.hoveredItem.row === hoveredItem.row;
-        if (!sameHovered || (!hoveredItem && state.hoveredItem)) {
-          state.hoveredItem = hoveredItem;
-          updateTooltip();
-        } else if (hoveredItem) {
-          state.hoveredItem = hoveredItem;
-          updateTooltip();
-        }
-      });
-      canvas.addEventListener('mouseleave', () => {
-        if (state.hoveredItem) {
-          state.hoveredItem = null;
-          updateTooltip();
-        }
-      });
-      canvas.addEventListener('mousedown', (event) => {
-        event.preventDefault();
-        state.dragging = true;
-        state.dragStartX = event.clientX;
-        state.dragViewStart = state.viewStart;
-        state.dragViewEnd = state.viewEnd;
-      });
-      window.addEventListener('mouseup', () => { state.dragging = false; });
-      canvas.addEventListener('click', (event) => {
-        const rect = canvas.getBoundingClientRect();
-        const canvasX = event.clientX - rect.left;
-        const canvasY = event.clientY - rect.top;
-        const rowIdx = rowIndexForCanvasY(canvasY);
-        if (canvasX < LABEL_WIDTH) {
-          if (rowIdx >= 0 && rowIdx < state.rows.length && state.viewMode !== 'expanded') {
-            const row = state.rows[rowIdx];
-            state.selectedItem = { row, item: rowSummaryItem(row), x: event.clientX, y: event.clientY };
-            if (state.expandedRows.has(row.row_id)) {
-              state.expandedRows.delete(row.row_id);
-            } else {
-              state.expandedRows.add(row.row_id);
-            }
-            layout();
-            updateTooltip();
-            scheduleDraw();
-          }
-          return;
-        }
-        state.cursorTime = xToTime(canvasX);
-        state.selectedItem = null;
-        if (rowIdx >= 0 && rowIdx < state.rows.length) {
-          const row = state.rows[rowIdx];
-          const item = findItemAt(row, rowIdx, canvasX, canvasY);
-          if (item) state.selectedItem = { row, item, x: event.clientX, y: event.clientY };
-        }
-        updateStats();
-        scheduleDraw();
-      });
-      canvas.addEventListener('wheel', (event) => {
-        const rect = canvas.getBoundingClientRect();
-        const x = event.clientX - rect.left;
-        if (x < LABEL_WIDTH) {
-          event.preventDefault();
-          viewport.scrollTop += event.deltaY;
-          return;
-        }
-        event.preventDefault();
-        if (Math.abs(event.deltaX) > Math.abs(event.deltaY) * 0.25) {
-          panTimelineByPixels(event.deltaX);
-          return;
-        }
-        handleZoom(x, event.deltaY);
-      }, { passive: false });
-      window.addEventListener('resize', () => {
-        layout();
-        scheduleDraw();
-      });
-    }
-
-    async function boot() {
-      const response = await fetch(CACHE_FILE);
-      const data = await response.json();
-      document.getElementById('title').textContent = data.pt_path.split('/').pop() + ' trace timeline';
-      state.rawRows = data.rows || [];
-      const attemptValues = [...new Set(
-        state.rawRows.flatMap(row => (row.items || []).map(item => item.attempt)).filter(v => v != null)
-      )].sort((a, b) => a - b);
-      document.getElementById('attemptValue').innerHTML =
-        ['<option value="all">all attempt segments</option>'].concat(
-          attemptValues.map(value => `<option value="${value}">attempt ${value} segments</option>`)
-        ).join('');
-      state.globalStart = data.global_start;
-      state.globalEnd = data.global_end;
-      state.cursorTime = state.globalStart;
-      fitToRows(state.rawRows);
-      attachEvents();
-      applyFilterAndSort();
-    }
-
-    boot().catch((error) => {
-      document.body.innerHTML = `<pre style="padding:16px">${String(error)}</pre>`;
-      console.error(error);
-    });
-  </script>
-</body>
-</html>
-"""
-
-
-def ensure_html(paths: TimelinePaths) -> None:
-    title = f"{paths.pt_path.name} trace timeline"
-    html = HTML_TEMPLATE.replace("__CACHE_FILE__", paths.cache_path.name).replace("__TITLE__", title)
-    with paths.html_path.open("w", encoding="utf-8") as handle:
-        handle.write(html)
-
-
-class QuietHandler(SimpleHTTPRequestHandler):
-    def log_message(self, format: str, *args: Any) -> None:
-        return
-
-
-def serve_directory(directory: Path, port: int) -> None:
-    handler = functools.partial(QuietHandler, directory=str(directory))
-    with socketserver.TCPServer(("0.0.0.0", port), handler) as httpd:
-        print(f"Serving http://127.0.0.1:{port}/")
-        print("Press Ctrl+C to stop.")
-        try:
-            httpd.serve_forever()
-        except KeyboardInterrupt:
-            pass
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description=__doc__)
-    parser.add_argument("pt_path", help="Path to rollout debug dump .pt file")
-    parser.add_argument("--rebuild", action="store_true", help="Rebuild cache even if it already exists")
-    parser.add_argument(
-        "--serve",
-        action=argparse.BooleanOptionalAction,
-        default=True,
-        help="Start a local static file server for the generated HTML",
-    )
-    parser.add_argument("--port", type=int, default=9999, help="Port for --serve")
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    pt_path = Path(args.pt_path).expanduser().resolve()
-    if not pt_path.exists():
-        raise SystemExit(f"pt file not found: {pt_path}")
-
-    paths = _timeline_paths(pt_path)
-    cache_data = ensure_cache(paths, rebuild=args.rebuild)
-    ensure_html(paths)
-
-    print(f"pt: {paths.pt_path}")
-    print(f"cache: {paths.cache_path}")
-    print(f"html: {paths.html_path}")
-    print(f"samples: {cache_data['sample_count']}")
-
-    if args.serve:
-        serve_directory(paths.html_path.parent, args.port)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/train.py b/train.py
index 2404a0bbd1..7e7f752b3a 100644
--- a/train.py
+++ b/train.py
@@ -2,7 +2,7 @@
 
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
-from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking, update_tracking_open_metrics
+from slime.utils.logging_utils import configure_logger, init_tracking
 from slime.utils.misc import should_run_periodic_action
 
 
@@ -16,17 +16,13 @@ def train(args):
     # need to initialize rollout manager first to calculate num_rollout
     rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"])
 
-    # Update primary W&B with SGLang metrics endpoint now that servers are up.
-    router_addr = ray.get(rollout_manager.get_metrics_router_addr.remote())
-    update_tracking_open_metrics(args, router_addr)
-
     # create the actor and critic models
     actor_model, critic_model = create_training_models(args, pgs, rollout_manager)
 
     if args.offload_rollout:
         ray.get(rollout_manager.onload_weights.remote())
 
-    # Always push actor weights to rollout once weights are loaded.
+    # always update weight first so that sglang has the loaded weights from training.
     actor_model.update_weights()
 
     if args.check_weight_update_equal:
@@ -39,18 +35,19 @@ def train(args):
     if args.num_rollout == 0 and args.eval_interval is not None:
         ray.get(rollout_manager.eval.remote(rollout_id=0))
 
-    def offload_train(actor_trains_this_step):
-        # Each model auto-offloads after train() when offload_train is set,
-        # so we only need clear_memory for the non-offload case.
-        if not args.offload_train:
-            if not args.use_critic or actor_trains_this_step:
-                actor_model.clear_memory()
+    def offload_train(rollout_id):
+        if args.offload_train:
+            if args.use_critic:
+                critic_model.offload()
+                if rollout_id >= args.num_critic_only_steps:
+                    actor_model.offload()
             else:
-                critic_model.clear_memory()
+                actor_model.offload()
+        else:
+            actor_model.clear_memory()
 
     def save(rollout_id):
-        actor_trains_this_step = (not args.use_critic) or rollout_id >= args.num_critic_only_steps
-        if actor_trains_this_step:
+        if (not args.use_critic) or (rollout_id >= args.num_critic_only_steps):
             actor_model.save_model(
                 rollout_id,
                 force_sync=rollout_id == args.num_rollout - 1,
@@ -64,6 +61,7 @@ def save(rollout_id):
             ray.get(rollout_manager.save.remote(rollout_id))
 
     # train loop.
+    # note that for async training, one can change the position of the sync operation(ray.get).
     for rollout_id in range(args.start_rollout_id, args.num_rollout):
         if args.eval_interval is not None and rollout_id == 0 and not args.skip_eval_before_train:
             ray.get(rollout_manager.eval.remote(rollout_id))
@@ -73,25 +71,21 @@ def save(rollout_id):
         if args.offload_rollout:
             ray.get(rollout_manager.offload.remote())
 
-        actor_trains_this_step = (not args.use_critic) or rollout_id >= args.num_critic_only_steps
-
         if args.use_critic:
-            value_refs = critic_model.async_train(rollout_id, rollout_data_ref)
-            if actor_trains_this_step:
-                ray.get(actor_model.async_train(rollout_id, rollout_data_ref, external_data=value_refs))
-            else:
-                ray.get(value_refs)
+            critic_train_handle = critic_model.async_train(rollout_id, rollout_data_ref)
+            if rollout_id >= args.num_critic_only_steps:
+                ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
+            ray.get(critic_train_handle)
         else:
             ray.get(actor_model.async_train(rollout_id, rollout_data_ref))
 
         if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout):
             save(rollout_id)
 
-        offload_train(actor_trains_this_step)
+        offload_train(rollout_id)
         if args.offload_rollout:
             ray.get(rollout_manager.onload_weights.remote())
         actor_model.update_weights()
-
         if args.offload_rollout:
             ray.get(rollout_manager.onload_kv.remote())
 
@@ -99,7 +93,6 @@ def save(rollout_id):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
     ray.get(rollout_manager.dispose.remote())
-    finish_tracking(args)
 
 
 if __name__ == "__main__":
diff --git a/train_async.py b/train_async.py
index 6960bd0558..fc6c17d58c 100644
--- a/train_async.py
+++ b/train_async.py
@@ -2,7 +2,7 @@
 
 from slime.ray.placement_group import create_placement_groups, create_rollout_manager, create_training_models
 from slime.utils.arguments import parse_args
-from slime.utils.logging_utils import configure_logger, finish_tracking, init_tracking, update_tracking_open_metrics
+from slime.utils.logging_utils import configure_logger, init_tracking
 from slime.utils.misc import should_run_periodic_action
 
 
@@ -18,14 +18,10 @@ def train(args):
     # need to initialize rollout manager first to calculate num_rollout
     rollout_manager, num_rollout_per_epoch = create_rollout_manager(args, pgs["rollout"])
 
-    # Update primary W&B with SGLang metrics endpoint now that servers are up.
-    router_addr = ray.get(rollout_manager.get_metrics_router_addr.remote())
-    update_tracking_open_metrics(args, router_addr)
-
     # create the actor and critic models
     actor_model, critic_model = create_training_models(args, pgs, rollout_manager)
 
-    # Always push actor weights to rollout once weights are loaded.
+    # always update weight first so that sglang has the loaded weights from training.
     actor_model.update_weights()
 
     if args.check_weight_update_equal:
@@ -43,21 +39,18 @@ def train(args):
             rollout_data_next_future = rollout_manager.generate.remote(rollout_id + 1)
 
         if args.use_critic:
-            actor_trains_this_step = rollout_id >= args.num_critic_only_steps
-            value_refs = critic_model.async_train(rollout_id, rollout_data_curr_ref)
-            if actor_trains_this_step:
-                ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref, external_data=value_refs))
-            else:
-                ray.get(value_refs)
+            critic_train_handle = critic_model.async_train(rollout_id, rollout_data_curr_ref)
+            if rollout_id >= args.num_critic_only_steps:
+                ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
+            ray.get(critic_train_handle)
         else:
             ray.get(actor_model.async_train(rollout_id, rollout_data_curr_ref))
 
         if should_run_periodic_action(rollout_id, args.save_interval, num_rollout_per_epoch, args.num_rollout):
-            if (not args.use_critic) or rollout_id >= args.num_critic_only_steps:
-                actor_model.save_model(
-                    rollout_id,
-                    force_sync=rollout_id == args.num_rollout - 1,
-                )
+            actor_model.save_model(
+                rollout_id,
+                force_sync=rollout_id == args.num_rollout - 1,
+            )
             if args.use_critic:
                 critic_model.save_model(
                     rollout_id,
@@ -76,7 +69,6 @@ def train(args):
             ray.get(rollout_manager.eval.remote(rollout_id))
 
     ray.get(rollout_manager.dispose.remote())
-    finish_tracking(args)
 
 
 if __name__ == "__main__":