[AMD] Speed up DSV4 MoE weight loading from mmap views (#32315)

2026-08-01 23:47:47 -07:00
parent 7e509f690e
commit 37be4e9247
6 changed files with 91 additions and 0 deletions
@@ -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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $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_summary.md )" >> $GITHUB_STEP_SUMMARY || true
@@ -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_summary.md )" >> $GITHUB_STEP_SUMMARY || true
@@ -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_summary.md )" >> $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_summary.md )" >> $GITHUB_STEP_SUMMARY || true
          if [[ "${{ github.event_name == 'schedule' || inputs.continue_on_error }}" == "true" ]]; then
@@ -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("")
@@ -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(
@@ -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"]))