minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Feat: Add prefill context parallel support for deepseek v4 unified kv attention (#27928)

Browse Source 
Co-authored-by: Thomas Wang <thomawan@amd.com>
This commit is contained in:
Wang, FangYuan
2026-06-16 02:00:51 -07:00
committed by GitHub
co-authored by Thomas Wang
parent 149fabcca7
commit a362ba9da3
3 changed files with 213 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py
+56 -5
View File
@@ -320,7 +320,7 @@ class DSV4AttnMetadata:
f"!= pre_global_len={pre_global_len} (must remain global for compressor write path)"
)
def init_flashmla_related(self):
def init_flashmla_related(self, is_prefill: bool = False):
# c4_sparse_topk is set from model_config.index_topk per-model
# (small model: 512, large model: 1024).
assert self.c4_sparse_topk in (512, 1024), (
@@ -342,6 +342,8 @@ class DSV4AttnMetadata:
device=self.c4_topk_lengths_clamp1.device,
)
self.c4_sparse_page_indices = _pad_last_dim(self.c4_sparse_page_indices)
if is_prefill:
self.c4_sparse_raw_indices = torch.empty_like(self.c4_sparse_page_indices)
self.c1_flashmla_metadata = _create_flashmla_metadata()
self.c4_flashmla_metadata = _create_flashmla_metadata()
self.c128_flashmla_metadata = _create_flashmla_metadata()
@@ -1187,6 +1189,49 @@ class DeepseekV4HipRadixBackend(
cu_q = core_attn_metadata.unified.pf_cu_q
final_pos = core_attn_metadata.unified.pf_final_pos
# DSA CP (round-robin/interleave): unified_pf_* are built over the GLOBAL
# token layout, but under CP each rank owns only 1/cp_size of the queries
# (q/positions are local) while kv was all-gathered to the full sequence.
# Slice the per-query fields to this rank's tokens so their length matches
# the local query count T; values stay global so each local query still
# attends over the full all-gathered KV.
from sglang.srt.layers.attention.dsa.utils import (
is_dsa_prefill_cp_round_robin_split,
)
# NOTE (AMD/HIP only): this whole DSA-CP prefill handling lives in the
# HIP backend (DeepseekV4HipRadixBackend, selected only when is_hip()).
# The NVIDIA path uses DeepseekV4AttnBackend and never reaches here, so
# these CP changes do not affect B200/H200 execution.
_cp_size = get_attention_cp_size()
_cp_active = (
_cp_size > 1
and is_dsa_prefill_cp_round_robin_split()
and kv.shape[0] == _cp_size * T
and state_slot.shape[0] != T
)
state_slot_full = state_slot
final_pos_full = final_pos
positions_full = positions
if _cp_active:
_sl = slice(get_attention_cp_rank(), None, _cp_size)
state_slot = state_slot[_sl].contiguous()
chunk_start = chunk_start[_sl].contiguous()
cu_q = cu_q[_sl].contiguous()
final_pos = final_pos[_sl].contiguous()
# positions for the local queries are this rank's round-robin global
# positions {r, r+cp, r+2cp, ...}; forward_batch.positions is the full
# (padded) global layout, so slice it the same way instead of taking
# the first T entries (which would be the wrong, sequential 0..T-1).
positions = forward_batch.positions.to(torch.int64)[_sl].contiguous()
# The SWA ring must hold the FULL window on EVERY rank (decode and
# later chunks read this rank's ring). kv was all-gathered to the full
# sequence, so write the full kv with full global positions/state_slot
# instead of only this rank's 1/cp_size tokens.
positions_full = forward_batch.positions.to(torch.int64)[
: state_slot_full.shape[0]
].contiguous()
kpre_i, kpre_p, kext_i, kext_p = runtime.build_prefill_indices(
compress_ratio=compress_ratio,
state_slot=state_slot,
@@ -1219,15 +1264,21 @@ class DeepseekV4HipRadixBackend(
# write this chunk's SWA K into the ring for future chunks / decode
# only the final-window tokens per request
if save_kv_cache:
n_real = state_slot.shape[0]
# Under CP, write the FULL all-gathered window so every rank's ring is
# complete (decode / later chunks read the local ring). Without CP this
# is just the local kv + local metadata as before.
_ring_state_slot = state_slot_full if _cp_active else state_slot
_ring_final_pos = final_pos_full if _cp_active else final_pos
_ring_positions = positions_full if _cp_active else positions
n_real = _ring_state_slot.shape[0]
runtime.store_swa_into_unified(
kv=kv[:n_real],
state_slot=state_slot,
positions=positions[:n_real],
state_slot=_ring_state_slot,
positions=_ring_positions[:n_real],
unified_kv=unified,
win=win,
ring_stride=ring_stride,
final_pos=final_pos,
final_pos=_ring_final_pos,
)
return o
python/sglang/srt/models/deepseek_v4.py
+13
View File
@@ -822,6 +822,19 @@ class MQALayer(nn.Module):
# unified_kv prefill: keep bf16 kv; the backend writes
# the ring AFTER attention (2-source path).
kv = self._compute_kv_bf16(x_linear, positions, qkv_a=qkv_a)
# HIP/ROCm-only: the unified_kv 2-source prefill path is exclusive
# to DeepseekV4HipRadixBackend. Guard with _is_hip so this CP
# all-gather never enters the NVIDIA (DeepseekV4AttnBackend) path.
if use_cp and _is_hip:
# unified_kv + DSA CP: the 2-source prefill path needs the
# FULL current-chunk KV (extend source + ring write), so
# all-gather the per-rank bf16 KV across the CP group.
kv = cp_all_gather_rerange_output(
kv.contiguous(),
self.cp_size,
forward_batch,
torch.cuda.current_stream(),
)
elif use_cp:
# NSA CP: keep bf16 kv around for the cross-rank all-gather, then
# write to the FlashMLA cache after gather.
test/registered/amd/test_deepseek_v4_pro_fp4_cp.py
+144
View File
@@ -0,0 +1,144 @@
"""MI35x DeepSeek-V4-Pro FP4 prefill context-parallel (CP) accuracy test (8-GPU).
Shares the launch conventions of test_deepseek_v4_pro_fp4.py (same 1.6T model,
same env, same long launch timeout) but enables prefill CP via
``--enable-prefill-cp --cp-strategy interleave`` over the unified_kv backend.
Registry: nightly-amd-8-gpu-mi35x-deepseek-v4-pro suite
"""
import os
import unittest
from types import SimpleNamespace
from sglang.srt.utils import kill_process_tree
from sglang.test.ci.ci_register import register_amd_ci
from sglang.test.run_eval import run_eval
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=5400, suite="nightly-amd-8-gpu-mi35x-deepseek-v4-pro", 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
# Common DeepSeek-V4 env vars, aligned with test_deepseek_v4_pro_fp4.py, except
# SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton which the prefill-CP path requires.
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": "unified_kv_triton",
"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": "false",
}
# FP4 variant (matches test_deepseek_v4_pro_fp4.py; V4-Pro also auto-detects it).
FP4_ENV_VARS = {
"SGLANG_DSV4_FP4_EXPERTS": "true",
}
class TestDeepseekV4ProFp4CPInterleave(CustomTestCase):
"""DeepSeek-V4-Pro FP4 unified_kv prefill CP, interleave (round-robin-split), tp=8."""
@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",
"--dp",
"1",
"--enable-prefill-cp",
"--cp-strategy",
"interleave",
"--disable-radix-cache",
"--attention-backend",
"dsv4",
"--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,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
base_url=self.base_url,
model=self.model,
eval_name="gsm8k",
api="completion",
max_tokens=512,
num_examples=1319,
num_threads=32,
num_shots=5,
)
metrics = run_eval(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_a_gsm8k (deepseek-v4-pro-fp4-cp-interleave)\n"
f'{metrics["score"]=:.3f}\n'
)
self.assertGreater(metrics["score"], 0.92)
if __name__ == "__main__":
unittest.main()