From 37be4e9247759a88f3d55579ecd238cafe928770 Mon Sep 17 00:00:00 2001 From: Bingxu Chen Date: Sun, 2 Aug 2026 14:47:47 +0800 Subject: [PATCH] [AMD] Speed up DSV4 MoE weight loading from mmap views (#32315) --- .../workflows/nightly-test-amd-rocm720.yml | 9 ++++ .github/workflows/nightly-test-amd.yml | 1 + .github/workflows/pr-test-amd-rocm720.yml | 2 + python/sglang/srt/environ.py | 3 ++ .../srt/layers/moe/fused_moe_triton/layer.py | 26 ++++++++++ .../moe/test_copy_weight_views_before_h2d.py | 50 +++++++++++++++++++ 6 files changed, 91 insertions(+) create mode 100644 test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py diff --git a/.github/workflows/nightly-test-amd-rocm720.yml b/.github/workflows/nightly-test-amd-rocm720.yml index 8af8e9838..d6dfe0ebe 100644 --- a/.github/workflows/nightly-test-amd-rocm720.yml +++ b/.github/workflows/nightly-test-amd-rocm720.yml @@ -493,6 +493,7 @@ jobs: timeout-minutes: 180 run: | bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-gpt-oss --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1317,6 +1318,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-flash --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1329,6 +1331,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-flash --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1366,6 +1369,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-pro --nightly --timeout-per-file 14400 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1378,6 +1382,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-pro --nightly --timeout-per-file 14400 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1415,6 +1420,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1427,6 +1433,7 @@ jobs: echo "## SGLANG_HACK_FLASHMLA_BACKEND=triton" >> github_summary.md bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e SGLANG_HACK_FLASHMLA_BACKEND=triton \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1536,6 +1543,7 @@ jobs: run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-8-gpu-qwen3-235b --nightly --timeout-per-file 3600 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true @@ -1610,6 +1618,7 @@ jobs: run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-qwen35 --nightly --timeout-per-file 3600 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true diff --git a/.github/workflows/nightly-test-amd.yml b/.github/workflows/nightly-test-amd.yml index da995a2c3..492f7b1c5 100644 --- a/.github/workflows/nightly-test-amd.yml +++ b/.github/workflows/nightly-test-amd.yml @@ -1464,6 +1464,7 @@ jobs: run: | > github_summary.md # Clear summary file bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ python3 run_suite.py --hw amd --suite nightly-amd-accuracy-8-gpu-qwen35 --nightly --timeout-per-file 3600 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true diff --git a/.github/workflows/pr-test-amd-rocm720.yml b/.github/workflows/pr-test-amd-rocm720.yml index eb92573cf..2662f92ad 100644 --- a/.github/workflows/pr-test-amd-rocm720.yml +++ b/.github/workflows/pr-test-amd-rocm720.yml @@ -1279,6 +1279,7 @@ jobs: bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ -e SGLANG_DSV4_ACCURACY_ONLY=1 \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-v4-flash --nightly --timeout-per-file 7200 ${{ (github.event_name == 'schedule' || inputs.continue_on_error) && '--continue-on-error' || '' }} || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} @@ -1330,6 +1331,7 @@ jobs: bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ -e SGLANG_DSV4_ACCURACY_ONLY=1 \ + -e SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D=1 \ python3 registered/amd/test_deepseek_v4_pro_fp4.py || TEST_EXIT_CODE=$? echo "$(> $GITHUB_STEP_SUMMARY || true if [[ "${{ github.event_name == 'schedule' || inputs.continue_on_error }}" == "true" ]]; then diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5d21605de..2dbc6a04f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1095,6 +1095,9 @@ class Envs: # Set False when using FP4-to-FP8 converted DeepSeek V4 checkpoint. SGLANG_DSV4_FP4_EXPERTS = EnvBool(True) SGLANG_DSV4_FP4_DEQUANT = EnvBool(False) + # Copy rank-local MoE slices into independent CPU storage before H2D when + # they reference a larger mmap-backed checkpoint storage. + SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D = EnvBool(False) # Default reasoning_effort for dsv4 chat encoder when request doesn't set it. # Accepts "", "max", "high" (empty string means unset); other values filtered to None. SGLANG_DSV4_REASONING_EFFORT = EnvStr("") 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 006f0db17..8755afc3a 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -94,6 +94,29 @@ _is_npu = is_npu() _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip +def _copy_weight_view_before_h2d(loaded_weight: torch.Tensor) -> torch.Tensor: + """Copy a CPU tensor view into independent contiguous storage.""" + if loaded_weight.device.type != "cpu": + return loaded_weight + tensor_bytes = loaded_weight.numel() * loaded_weight.element_size() + needs_copy = not ( + loaded_weight.is_contiguous() + and loaded_weight.storage_offset() == 0 + and loaded_weight.untyped_storage().nbytes() == tensor_bytes + ) + if not needs_copy: + return loaded_weight + return loaded_weight.clone(memory_format=torch.contiguous_format) + + +def _maybe_copy_weight_view_before_h2d( + loaded_weight: torch.Tensor, +) -> torch.Tensor: + if not envs.SGLANG_MOE_COPY_WEIGHT_VIEWS_BEFORE_H2D.get(): + return loaded_weight + return _copy_weight_view_before_h2d(loaded_weight) + + def _get_deepep_comm_group(a2a_backend): group = get_tp_group().device_group @@ -559,6 +582,7 @@ class FusedMoE(torch.nn.Module): ): # for per channel weight quantization if shard_id == "w2": + loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight) expert_data.copy_(loaded_weight) elif shard_id in ("w1", "w3"): self._load_w13( @@ -631,6 +655,7 @@ class FusedMoE(torch.nn.Module): ) expert_data = expert_data.narrow(shard_dim, start, shard_size) + loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight) expert_data.copy_(loaded_weight) def _load_w2( @@ -701,6 +726,7 @@ class FusedMoE(torch.nn.Module): ) # w2, down_proj: Load into only logical weight of w2. + loaded_weight = _maybe_copy_weight_view_before_h2d(loaded_weight) expert_data.copy_(loaded_weight) def _maybe_load_fp8_shared_expert_as_fp4( diff --git a/test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py b/test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py new file mode 100644 index 000000000..7d59e0bb2 --- /dev/null +++ b/test/registered/unit/layers/moe/test_copy_weight_views_before_h2d.py @@ -0,0 +1,50 @@ +import torch + +from sglang.srt.layers.moe.fused_moe_triton.layer import ( + _copy_weight_view_before_h2d, +) +from sglang.test.ci.ci_register import register_cpu_ci + +register_cpu_ci(est_time=2, suite="base-a-test-cpu") + + +def _assert_independent_copy(source: torch.Tensor, result: torch.Tensor) -> None: + assert torch.equal(result, source) + assert result is not source + assert result.is_contiguous() + assert result.storage_offset() == 0 + assert result.untyped_storage().nbytes() == result.numel() * result.element_size() + + +def test_skips_copy_for_exact_storage(): + source = torch.arange(16).reshape(4, 4) + assert _copy_weight_view_before_h2d(source) is source + + +def test_copies_zero_offset_storage_view(): + backing = torch.arange(32).reshape(8, 4) + source = backing.narrow(0, 0, 2) + assert source.storage_offset() == 0 + assert source.untyped_storage().nbytes() > source.numel() * source.element_size() + _assert_independent_copy(source, _copy_weight_view_before_h2d(source)) + + +def test_copies_nonzero_offset_storage_view(): + backing = torch.arange(32).reshape(8, 4) + source = backing.narrow(0, 6, 2) + assert source.storage_offset() != 0 + _assert_independent_copy(source, _copy_weight_view_before_h2d(source)) + + +def test_copies_noncontiguous_tensor(): + source = torch.arange(16).reshape(4, 4).transpose(0, 1) + assert not source.is_contiguous() + _assert_independent_copy(source, _copy_weight_view_before_h2d(source)) + + +if __name__ == "__main__": + import sys + + import pytest + + sys.exit(pytest.main([__file__, "-v"]))