From cfb7fb4fad032687306d357f7dbdd4b1c67e1db9 Mon Sep 17 00:00:00 2001 From: Clint Date: Wed, 3 Jun 2026 18:56:05 -0500 Subject: [PATCH] [AMD] Fix TP2 DeepSeek-R1 nhead=64 MLA decode crash and add nightly coverage (#27188) Co-authored-by: clintg6 <7388379+clintg6@users.noreply.github.com> --- .../workflows/nightly-test-amd-rocm720.yml | 73 +++++++ .github/workflows/nightly-test-amd.yml | 73 +++++++ .../srt/layers/attention/aiter_backend.py | 2 +- .../mi35x/test_deepseek_r1_mxfp4_tp2_mi35x.py | 185 ++++++++++++++++++ .../mi35x/test_deepseek_r1_mxfp4_tp4_mi35x.py | 184 +++++++++++++++++ test/run_suite.py | 2 + 6 files changed, 518 insertions(+), 1 deletion(-) create mode 100644 test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp2_mi35x.py create mode 100644 test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp4_mi35x.py diff --git a/.github/workflows/nightly-test-amd-rocm720.yml b/.github/workflows/nightly-test-amd-rocm720.yml index a59532244..fcd3da890 100644 --- a/.github/workflows/nightly-test-amd-rocm720.yml +++ b/.github/workflows/nightly-test-amd-rocm720.yml @@ -38,6 +38,8 @@ on: - nightly-4-gpu-rocm720 # 2-GPU GLM-5.1-MXFP4 (MI35x only) - nightly-2-gpu-mi35x-glm51-mxfp4-rocm720 + # 2-GPU DeepSeek-R1-MXFP4 TP2 (MI35x only) + - nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2-rocm720 # 8-GPU GPT-OSS (MI30x mixes Grok1-FP8; MI35x mixes Qwen3-Coder-Next) - nightly-accuracy-8-gpu-rocm720 - nightly-accuracy-8-gpu-mi35x-rocm720 @@ -61,6 +63,7 @@ on: - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720 + - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-hicache-rocm720 # 8-GPU DeepSeek-V4 (MI35x only) - nightly-8-gpu-mi35x-deepseek-v4-flash-rocm720 @@ -262,6 +265,40 @@ jobs: echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} + nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2-rocm720: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2-rocm720,')) + runs-on: linux-mi35x-gpu-2 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ci/amd/ensure_vram_clear.sh rocm + + - name: Setup docker (ROCm 7.2) + run: | + touch github_summary.md + bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps + bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate + + - name: Accuracy Test MI35x ROCm 7.2 (2-GPU DeepSeek-R1-MXFP4 TP2) + timeout-minutes: 180 + run: | + > github_summary.md + bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ + python3 run_suite.py --hw amd --suite nightly-amd-2-gpu-mi35x-deepseek-r1-mxfp4-tp2 --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} + nightly-accuracy-2-gpu-vlm-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu-vlm-rocm720,')) runs-on: linux-mi325-2gpu-sglang @@ -1023,6 +1060,40 @@ jobs: echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} + nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4-rocm720: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4-rocm720,')) + runs-on: linux-mi35x-gpu-8 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ci/amd/ensure_vram_clear.sh rocm + + - name: Setup docker (ROCm 7.2) + run: | + touch github_summary.md + bash scripts/ci/amd/amd_ci_start_container.sh --rocm-version rocm720 + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh --skip-test-time-deps + bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate + + - name: Accuracy Test MI35x ROCm 7.2 (DeepSeek-R1-MXFP4 TP4) + timeout-minutes: 180 + run: | + > github_summary.md + bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ + python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-tp4 --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} + nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720,')) runs-on: linux-mi35x-gpu-8 @@ -1717,6 +1788,7 @@ jobs: - nightly-perf-2-gpu-vlm-rocm720 - nightly-4-gpu-rocm720 - nightly-2-gpu-mi35x-glm51-mxfp4-rocm720 + - nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2-rocm720 # 8-GPU GPT-OSS (MI30x mixes Grok1-FP8; MI35x mixes Qwen3-Coder-Next) - nightly-accuracy-8-gpu-rocm720 - nightly-accuracy-8-gpu-mi35x-rocm720 @@ -1740,6 +1812,7 @@ jobs: - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion-rocm720 + - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4-rocm720 - nightly-8-gpu-mi35x-deepseek-r1-hicache-rocm720 # 8-GPU DeepSeek-V4 (MI35x only) - nightly-8-gpu-mi35x-deepseek-v4-flash-rocm720 diff --git a/.github/workflows/nightly-test-amd.yml b/.github/workflows/nightly-test-amd.yml index 8cde330de..56d9a41f7 100644 --- a/.github/workflows/nightly-test-amd.yml +++ b/.github/workflows/nightly-test-amd.yml @@ -38,6 +38,8 @@ on: - nightly-4-gpu # 2-GPU GLM-5.1-MXFP4 (MI35x only) - nightly-2-gpu-mi35x-glm51-mxfp4 + # 2-GPU DeepSeek-R1-MXFP4 TP2 (MI35x only) + - nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2 # 8-GPU GPT-OSS (MI30x mixes Grok1-FP8; MI35x mixes Qwen3-Coder-Next) - nightly-accuracy-8-gpu - nightly-accuracy-8-gpu-mi35x @@ -61,6 +63,7 @@ on: - nightly-8-gpu-mi35x-deepseek-r1-mxfp4 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion + - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4 - nightly-8-gpu-mi35x-deepseek-r1-hicache # 8-GPU Kimi-K2.6 (MI30x + MI35x) - nightly-8-gpu-kimi-k26 @@ -264,6 +267,40 @@ jobs: echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} + nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2,')) + runs-on: linux-mi35x-gpu-2 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ci/amd/ensure_vram_clear.sh rocm + + - name: Setup docker + run: | + touch github_summary.md + bash scripts/ci/amd/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh + bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate + + - name: Accuracy Test MI35x (2-GPU DeepSeek-R1-MXFP4 TP2) + timeout-minutes: 180 + run: | + > github_summary.md + bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ + python3 run_suite.py --hw amd --suite nightly-amd-2-gpu-mi35x-deepseek-r1-mxfp4-tp2 --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} + nightly-accuracy-2-gpu-vlm: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-accuracy-2-gpu-vlm,')) runs-on: linux-mi325-2gpu-sglang @@ -1028,6 +1065,40 @@ jobs: echo "$(> $GITHUB_STEP_SUMMARY || true exit ${TEST_EXIT_CODE:-0} + nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4: + if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4,')) + runs-on: linux-mi35x-gpu-8 + steps: + - name: Checkout code + uses: actions/checkout@v4 + with: + ref: ${{ inputs.ref || github.sha }} + + - name: Ensure VRAM is clear + run: bash scripts/ci/amd/ensure_vram_clear.sh rocm + + - name: Setup docker + run: | + touch github_summary.md + bash scripts/ci/amd/amd_ci_start_container.sh + env: + GITHUB_WORKSPACE: ${{ github.workspace }} + + - name: Install dependencies + run: | + bash scripts/ci/amd/amd_ci_install_dependency.sh + bash scripts/ci/amd/amd_ci_exec.sh pip install tabulate + + - name: Accuracy Test MI35x (DeepSeek-R1-MXFP4 TP4) + timeout-minutes: 180 + run: | + > github_summary.md + bash scripts/ci/amd/amd_ci_exec.sh -w /sglang-checkout/test \ + -e GITHUB_STEP_SUMMARY="/sglang-checkout/github_summary.md" \ + python3 run_suite.py --hw amd --suite nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-tp4 --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} + nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8: if: (github.repository == 'sgl-project/sglang' || github.event_name == 'pull_request') && (!(inputs.job_filter || inputs.job_select) || (inputs.job_filter || inputs.job_select) == 'all' || contains(format(',{0},', inputs.job_filter || inputs.job_select), ',nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8,')) runs-on: linux-mi35x-gpu-8 @@ -1648,6 +1719,7 @@ jobs: # - nightly-perf-2-gpu-vlm # excluded: perf failures don't block CI - nightly-4-gpu - nightly-2-gpu-mi35x-glm51-mxfp4 + - nightly-2-gpu-mi35x-deepseek-r1-mxfp4-tp2 # 8-GPU GPT-OSS (MI30x mixes Grok1-FP8; MI35x mixes Qwen3-Coder-Next) - nightly-accuracy-8-gpu - nightly-accuracy-8-gpu-mi35x @@ -1671,6 +1743,7 @@ jobs: - nightly-8-gpu-mi35x-deepseek-r1-mxfp4 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-kv-fp8 - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-ar-fusion + - nightly-8-gpu-mi35x-deepseek-r1-mxfp4-tp4 - nightly-8-gpu-mi35x-deepseek-r1-hicache # 8-GPU Kimi-K2.6 (MI30x + MI35x) - nightly-8-gpu-kimi-k26 diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 982b98677..692ac2728 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -281,7 +281,7 @@ class AiterAttnBackend(AttentionBackend): # so all num_head size does not use qh16 kernel to simulate # it should not use fake-nps (fast_mode = False, intra_batch_mode = True) # it will cause gpu-fault or accuracy issue - if self.num_head == 32 or self.num_head == 128: + if self.num_head in (32, 64, 128): fast_mode = True intra_batch_mode = False diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp2_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp2_mi35x.py new file mode 100644 index 000000000..1933579dc --- /dev/null +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp2_mi35x.py @@ -0,0 +1,185 @@ +"""MI35x DeepSeek-R1-MXFP4 TP=2 GSM8K AITER MLA regression. + +DeepSeek-R1 has 128 attention heads, so TP=2 gives 64 heads per rank. This +covers the AITER persistent MLA decode metadata path that previously crashed +with GPU memory access faults. + +Registry: nightly-amd-2-gpu-mi35x-deepseek-r1-mxfp4-tp2 suite +""" + +import ast +import os + +# Set HF cache for MI35x. +os.environ.setdefault("HF_HOME", "/data2/models/huggingface") +os.environ.setdefault("HF_HUB_CACHE", "/data2/models/huggingface/hub") + +import re +import time +import unittest +from typing import Tuple + +import numpy as np + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) +from sglang.utils import download_and_cache_file, read_jsonl + +register_amd_ci( + est_time=1800, + suite="nightly-amd-2-gpu-mi35x-deepseek-r1-mxfp4-tp2", + nightly=True, +) + +INVALID = -9999999 + +DEEPSEEK_R1_MXFP4_LOCAL_PATH = "/data2/models/amd-DeepSeek-R1-MXFP4-Preview" +DEEPSEEK_R1_MXFP4_HF_MODEL_ID = "amd/DeepSeek-R1-MXFP4-Preview" +SERVER_LAUNCH_TIMEOUT = 3600 +GSM8K_ACCURACY_THRESHOLD = 0.93 + + +def get_model_path() -> str: + env_path = os.environ.get("DEEPSEEK_R1_MXFP4_MODEL_PATH") + if env_path: + return env_path + if os.path.exists(DEEPSEEK_R1_MXFP4_LOCAL_PATH): + return DEEPSEEK_R1_MXFP4_LOCAL_PATH + return DEEPSEEK_R1_MXFP4_HF_MODEL_ID + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def run_gsm8k_benchmark( + base_url: str, + num_questions: int = 200, + num_shots: int = 5, + parallel: int = 64, +) -> Tuple[float, float, float]: + import sglang as sgl + from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + @sgl.function + def few_shot_gsm8k(s, question): + s += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + backend = RuntimeEndpoint(base_url) + sgl.set_default_backend(backend) + + tic = time.perf_counter() + states = few_shot_gsm8k.run_batch( + arguments, temperature=0, num_threads=parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + preds = [get_answer_value(states[i]["answer"]) for i in range(len(states))] + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + return float(acc), float(invalid), float(latency) + + +class TestDeepSeekR1MXFP4TP2MI35x(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.model = get_model_path() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "200")) + + env = os.environ.copy() + env["SGLANG_USE_AITER"] = "1" + env["SGLANG_AITER_MLA_PERSIST"] = "1" + + cls.process = popen_launch_server( + model=cls.model, + base_url=cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--attention-backend", + "aiter", + "--tp", + "2", + "--chunked-prefill-size", + "131072", + "--disable-radix-cache", + "--mem-fraction-static", + "0.85", + "--trust-remote-code", + "--kv-cache-dtype", + "fp8_e4m3", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ], + env=env, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + acc, invalid, latency = run_gsm8k_benchmark( + self.base_url, num_questions=self.num_questions + ) + print(f"accuracy={acc:.3f} invalid={invalid:.3f} latency={latency:.1f}s") + + if is_in_ci(): + write_github_step_summary( + "### DeepSeek-R1-MXFP4 TP=2 GSM8K (MI35x)\n\n" + "| Model | TP | Examples | Accuracy | Invalid | Threshold | Latency |\n" + "| ----- | -- | -------- | -------- | ------- | --------- | ------- |\n" + f"| {self.model} | 2 | {self.num_questions} | {acc:.3f} | " + f"{invalid:.3f} | {GSM8K_ACCURACY_THRESHOLD:.2f} | {latency:.1f}s |\n" + ) + + self.assertGreaterEqual(acc, GSM8K_ACCURACY_THRESHOLD) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp4_mi35x.py b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp4_mi35x.py new file mode 100644 index 000000000..2e8f2858d --- /dev/null +++ b/test/registered/amd/accuracy/mi35x/test_deepseek_r1_mxfp4_tp4_mi35x.py @@ -0,0 +1,184 @@ +"""MI35x DeepSeek-R1-MXFP4 TP=4 GSM8K AITER MLA regression. + +DeepSeek-R1 has 128 attention heads, so TP=4 gives 32 heads per rank. This +covers the AITER persistent MLA decode metadata path for the nhead=32 case. + +Registry: nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-tp4 suite +""" + +import ast +import os + +# Set HF cache for MI35x. +os.environ.setdefault("HF_HOME", "/data2/models/huggingface") +os.environ.setdefault("HF_HUB_CACHE", "/data2/models/huggingface/hub") + +import re +import time +import unittest +from typing import Tuple + +import numpy as np + +from sglang.srt.utils import kill_process_tree +from sglang.test.ci.ci_register import register_amd_ci +from sglang.test.test_utils import ( + DEFAULT_URL_FOR_TEST, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) +from sglang.utils import download_and_cache_file, read_jsonl + +register_amd_ci( + est_time=1800, + suite="nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-tp4", + nightly=True, +) + +INVALID = -9999999 + +DEEPSEEK_R1_MXFP4_LOCAL_PATH = "/data2/models/amd-DeepSeek-R1-MXFP4-Preview" +DEEPSEEK_R1_MXFP4_HF_MODEL_ID = "amd/DeepSeek-R1-MXFP4-Preview" +SERVER_LAUNCH_TIMEOUT = 3600 +GSM8K_ACCURACY_THRESHOLD = 0.93 + + +def get_model_path() -> str: + env_path = os.environ.get("DEEPSEEK_R1_MXFP4_MODEL_PATH") + if env_path: + return env_path + if os.path.exists(DEEPSEEK_R1_MXFP4_LOCAL_PATH): + return DEEPSEEK_R1_MXFP4_LOCAL_PATH + return DEEPSEEK_R1_MXFP4_HF_MODEL_ID + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def run_gsm8k_benchmark( + base_url: str, + num_questions: int = 200, + num_shots: int = 5, + parallel: int = 64, +) -> Tuple[float, float, float]: + import sglang as sgl + from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + @sgl.function + def few_shot_gsm8k(s, question): + s += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + backend = RuntimeEndpoint(base_url) + sgl.set_default_backend(backend) + + tic = time.perf_counter() + states = few_shot_gsm8k.run_batch( + arguments, temperature=0, num_threads=parallel, progress_bar=True + ) + latency = time.perf_counter() - tic + + preds = [get_answer_value(states[i]["answer"]) for i in range(len(states))] + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + return float(acc), float(invalid), float(latency) + + +class TestDeepSeekR1MXFP4TP4MI35x(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.model = get_model_path() + cls.base_url = DEFAULT_URL_FOR_TEST + cls.num_questions = int(os.environ.get("GSM8K_NUM_QUESTIONS", "200")) + + env = os.environ.copy() + env["SGLANG_USE_AITER"] = "1" + env["SGLANG_AITER_MLA_PERSIST"] = "1" + + cls.process = popen_launch_server( + model=cls.model, + base_url=cls.base_url, + timeout=SERVER_LAUNCH_TIMEOUT, + other_args=[ + "--attention-backend", + "aiter", + "--tp", + "4", + "--chunked-prefill-size", + "131072", + "--disable-radix-cache", + "--mem-fraction-static", + "0.85", + "--trust-remote-code", + "--kv-cache-dtype", + "fp8_e4m3", + "--model-loader-extra-config", + '{"enable_multithread_load": true}', + ], + env=env, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + def test_gsm8k(self): + acc, invalid, latency = run_gsm8k_benchmark( + self.base_url, num_questions=self.num_questions + ) + print(f"accuracy={acc:.3f} invalid={invalid:.3f} latency={latency:.1f}s") + + if is_in_ci(): + write_github_step_summary( + "### DeepSeek-R1-MXFP4 TP=4 GSM8K (MI35x)\n\n" + "| Model | TP | Examples | Accuracy | Invalid | Threshold | Latency |\n" + "| ----- | -- | -------- | -------- | ------- | --------- | ------- |\n" + f"| {self.model} | 4 | {self.num_questions} | {acc:.3f} | " + f"{invalid:.3f} | {GSM8K_ACCURACY_THRESHOLD:.2f} | {latency:.1f}s |\n" + ) + + self.assertGreaterEqual(acc, GSM8K_ACCURACY_THRESHOLD) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/run_suite.py b/test/run_suite.py index 74f03d40d..551a36645 100644 --- a/test/run_suite.py +++ b/test/run_suite.py @@ -114,6 +114,8 @@ NIGHTLY_SUITES = { "nightly-amd-1-gpu", "nightly-amd-1-gpu-mi35x", "nightly-amd-1-gpu-zimage-turbo", + "nightly-amd-2-gpu-mi35x-deepseek-r1-mxfp4-tp2", + "nightly-amd-8-gpu-mi35x-deepseek-r1-mxfp4-tp4", "nightly-amd-4-gpu", "nightly-amd-8-gpu", "nightly-amd-vlm",