diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py index 91d16a20e..e31ac5775 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 5b46fb55b..765eeb92a 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -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. diff --git a/test/registered/amd/test_deepseek_v4_pro_fp4_cp.py b/test/registered/amd/test_deepseek_v4_pro_fp4_cp.py new file mode 100644 index 000000000..5a2732f71 --- /dev/null +++ b/test/registered/amd/test_deepseek_v4_pro_fp4_cp.py @@ -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()