[AMD] Fix deepseek-v4 mtp accept length issue (#28520)

Co-authored-by: Bingxu Chen <bingxche@amd.com>
This commit is contained in:
Thomas Wang
2026-06-17 11:13:55 -07:00
committed by GitHub
co-authored by Bingxu Chen
parent 8aaca72c21
commit f5b041622b
3 changed files with 249 additions and 7 deletions
@@ -68,6 +68,7 @@ on:
# 8-GPU DeepSeek-V4 (MI35x only)
- nightly-8-gpu-mi35x-deepseek-v4-flash-rocm720
- nightly-8-gpu-mi35x-deepseek-v4-pro-rocm720
- nightly-8-gpu-mi35x-deepseek-v4-pro-mtp-rocm720
# 8-GPU Kimi-K2.6 (MI30x + MI35x)
- nightly-8-gpu-kimi-k26-rocm720
- nightly-8-gpu-mi35x-kimi-k26-rocm720
@@ -1321,6 +1322,54 @@ jobs:
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
nightly-8-gpu-mi35x-deepseek-v4-pro-mtp-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-v4-pro-mtp-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: |
# --skip-test-time-deps: GSM8K + send_one don't need lmms-eval / human-eval.
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 + MTP Accept Length Test MI35x ROCm 7.2 (8-GPU DeepSeek-V4-Pro FP8 MTP, unified_kv_triton)
timeout-minutes: 240
run: |
> github_summary.md # Clear summary file
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 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
- name: Accuracy + MTP Accept Length Test MI35x ROCm 7.2 (8-GPU DeepSeek-V4-Pro FP8 MTP, triton)
if: ${{ !cancelled() }}
timeout-minutes: 240
run: |
> github_summary.md # Clear summary file
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 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
# ==============================================================================
# 8-GPU Kimi-K2.6 (MI30x + MI35x)
# ==============================================================================
@@ -1879,6 +1928,7 @@ jobs:
# 8-GPU DeepSeek-V4 (MI35x only)
- nightly-8-gpu-mi35x-deepseek-v4-flash-rocm720
- nightly-8-gpu-mi35x-deepseek-v4-pro-rocm720
- nightly-8-gpu-mi35x-deepseek-v4-pro-mtp-rocm720
# 8-GPU Kimi-K2.6 (MI30x + MI35x)
- nightly-8-gpu-kimi-k26-rocm720
- nightly-8-gpu-mi35x-kimi-k26-rocm720
@@ -1316,24 +1316,37 @@ class DeepseekV4HipRadixBackend(
recompute at store time, matching the pre-cache per-layer behavior, for
paths that never ran the decode-stream init (eager prefill/extend, idle,
or a batch re-padded after init -> shape mismatch).
Cached swa_loc is computed once from committed positions, so every draft-decode
step would reuse the same ring slot and break the chain. Recompute from the live
per-step positions; only the draft path is affected, the rest keeps the fast path.
"""
positions = forward_batch.positions
core = getattr(self.forward_metadata, "core_attn_metadata", None)
unified = getattr(core, "unified", None) if core is not None else None
cached = unified.swa_loc if unified is not None else None
is_multistep_draft_decode = (
forward_batch.forward_mode.is_decode_or_idle()
and self.speculative_num_steps > 1
)
if (
cached is not None
and not forward_batch.forward_mode.is_idle()
and cached.shape[0] == positions.shape[0]
and not is_multistep_draft_decode
):
return cached
ring = self.token_to_kv_pool.unified_swa_ring_size
req_slot = forward_batch.req_pool_indices.to(torch.int64)
if req_slot.shape[0] != positions.shape[0]:
req_slot = req_slot.repeat_interleave(
positions.shape[0] // req_slot.shape[0]
result = cached
else:
ring = self.token_to_kv_pool.unified_swa_ring_size
req_slot = forward_batch.req_pool_indices.to(torch.int64)
if req_slot.shape[0] != positions.shape[0]:
req_slot = req_slot.repeat_interleave(
positions.shape[0] // req_slot.shape[0]
)
result = (req_slot * ring + positions.to(torch.int64) % ring).to(
torch.int32
)
return (req_slot * ring + positions.to(torch.int64) % ring).to(torch.int32)
return result
def store_cache(
self, layer_id: int, swa_k: torch.Tensor, forward_batch: ForwardBatch
@@ -0,0 +1,179 @@
"""MI35x DeepSeek-V4-Pro FP4 + MTP Test (8-GPU)
- Accuracy: GSM8K few-shot eval
- Acceptance: mtp acc length eval
Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp suite
"""
import os
import unittest
from types import SimpleNamespace
import requests
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k
from sglang.test.send_one import BenchArgs, send_one_prompt
from sglang.test.test_utils import (
DEFAULT_URL_FOR_TEST,
CustomTestCase,
is_in_ci,
popen_launch_server,
write_github_step_summary,
)
register_amd_ci(
est_time=7200, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro-mtp", nightly=True
)
DEEPSEEK_V4_PRO_FP4_MODEL_PATH = os.environ.get(
"DEEPSEEK_V4_PRO_MODEL_PATH_FP4", "deepseek-ai/DeepSeek-V4-Pro"
)
# Pro is 1.6T; weight load + warmup is much longer than Flash 285B.
SERVER_LAUNCH_TIMEOUT = 5400
FLASHMLA_BACKEND = os.environ.get("SGLANG_HACK_FLASHMLA_BACKEND", "unified_kv_triton")
GSM8K_ACCURACY_THRESHOLD = 0.92
AVG_SPEC_ACCEPT_LENGTH_THRESHOLD = 2.8
# Common DeepSeek-V4 env vars (AMD ROCm 7.2 path: AITER indexer + triton attn + ROCm700A).
COMMON_ENV_VARS = {
"SGLANG_DEFAULT_THINKING": "1",
"SGLANG_DSV4_REASONING_EFFORT": "max",
"SGLANG_OPT_DEEPGEMM_HC_PRENORM": "false",
"SGLANG_USE_AITER": "1",
"SGLANG_USE_ROCM700A": "1",
"SGLANG_OPT_USE_FUSED_COMPRESS": "true",
"SGLANG_OPT_USE_FUSED_COMPRESS_TRITON": "true",
"SGLANG_HACK_FLASHMLA_BACKEND": FLASHMLA_BACKEND,
"SGLANG_OPT_FP8_WO_A_GEMM": "false",
"SGLANG_OPT_USE_JIT_INDEXER_METADATA": "false",
"SGLANG_OPT_USE_TOPK_V2": "false",
"SGLANG_OPT_USE_AITER_INDEXER": "true",
"SGLANG_OPT_USE_TILELANG_INDEXER": "false",
"SGLANG_OPT_USE_TILELANG_MHC_PRE": "false",
"SGLANG_OPT_USE_TILELANG_MHC_POST": "false",
"SGLANG_FP8_PAGED_MQA_LOGITS_TORCH": "1",
"SGLANG_OPT_USE_MULTI_STREAM_OVERLAP": "false",
"SGLANG_ROCM_USE_MULTI_STREAM": "false",
"AITER_BF16_FP8_MOE_BOUND": "0",
"SGLANG_EAGER_INPUT_NO_COPY": "1",
}
FP4_ENV_VARS = {
"SGLANG_DSV4_FP4_EXPERTS": "true",
}
class TestDeepseekV4ProFp4MTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V4_PRO_FP4_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
env = os.environ.copy()
env.update(COMMON_ENV_VARS)
env.update(FP4_ENV_VARS)
other_args = [
"--trust-remote-code",
"--tp",
"8",
"--disable-radix-cache",
"--attention-backend",
"dsv4",
# MTP / EAGLE speculative decoding (NextN head from the base model).
"--speculative-algorithm",
"EAGLE",
"--speculative-num-steps",
"3",
"--speculative-eagle-topk",
"1",
"--speculative-num-draft-tokens",
"4",
"--max-running-requests",
"256",
"--page-size",
"256",
"--mem-fraction-static",
"0.90",
"--swa-full-tokens-ratio",
"0.1",
"--chunked-prefill-size",
"8192",
"--disable-shared-experts-fusion",
"--tool-call-parser",
"deepseekv4",
"--reasoning-parser",
"deepseek-v4",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=SERVER_LAUNCH_TIMEOUT,
other_args=other_args,
env=env,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(self):
# `a` prefix to run first (alphabetically) and warm up the server.
requests.get(self.base_url + "/flush_cache")
args = SimpleNamespace(
num_shots=8,
data_path=None,
num_questions=1319,
parallel=1319,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
server_info = requests.get(self.base_url + "/server_info")
avg_spec_accept_length = server_info.json()["internal_states"][0][
"avg_spec_accept_length"
]
print(f"{avg_spec_accept_length=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v4-pro-fp4 MTP, {FLASHMLA_BACKEND})\n"
f'{metrics["accuracy"]=:.3f}\n'
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], GSM8K_ACCURACY_THRESHOLD)
self.assertGreater(avg_spec_accept_length, AVG_SPEC_ACCEPT_LENGTH_THRESHOLD)
def test_b_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
acc_length, speed = send_one_prompt(args)
print(f"{acc_length=:.2f} {speed=:.2f}")
if is_in_ci():
write_github_step_summary(
f"### test_bs_1_speed (deepseek-v4-pro-fp4 MTP, {FLASHMLA_BACKEND})\n"
f"{acc_length=:.2f}\n"
f"{speed=:.2f} token/s\n"
)
self.assertGreater(acc_length, AVG_SPEC_ACCEPT_LENGTH_THRESHOLD)
if __name__ == "__main__":
# run_suite.py's run_one_file launches each test file with `python3 <file> -f`,
# which enables unittest fail-fast. test_a_gsm8k (accuracy + accept length) and
# test_b_bs_1_speed are independent measurements sharing one expensive server
# launch; strip `-f` so later methods still run if an earlier one fails.
import sys
sys.argv = [a for a in sys.argv if a not in ("-f", "--failfast")]
unittest.main()