Repository navigation
[Feat] [AMD] dsv4: Support interleave prefill CP with breakable CUDA graphs - #43819
Open
AMD-yanfeiwang wants to merge 3 commits into
Open
AMD-yanfeiwang wants to merge 3 commits into
AMD-yanfeiwang wants to merge 3 commits into
Conversation
2 of 3 tasks
AMD-yanfeiwang
marked this pull request as ready for review
October 13, 2026 13:18
AMD-yanfeiwang
requested review from
BBuf,
Edwardf0t1,
Fridge003,
HaiShaw,
Qiaolin-Yu,
Ying1123,
ch-wan,
hebiao064,
hnyls2002,
ispobock and
merrymercy
as code owners
October 13, 2026 13:18
…phs on ROCm Port sgl-project#38603 to the HIP DSv4 attention backend: - HIP init_forward_metadata_prefill writes the global KV fields with the bucket-sized out_cache_loc under prefill graphs (num_write_tokens), passes the CP local_index and pads swa_replay_start, matching the CUDA backend. - CP BCG keeps forward_batch.positions global and feeds the captured body CP-local positions and ids through _cp_positions/_cp_input_ids, so the HIP unified-KV SWA ring store sees global positions. - CP BCG sets input_ids_global (rank-major for TP MoE, local for a2a) so the DSv4 hash-routed MoE layers route the right tokens (from sgl-project#40333).
… tail DSV4's body returns (hidden_states, pre_hc_head), and both CP tails (eager and breakable-graph replay) all-gathered both. The logits processor reads pre_hc_head only to store hidden states for MTP/EAGLE, and DSpark drops it, so plain serving and DSpark gathered ~4x the hidden-state bytes for nothing. On large prefill buckets that gather also exceeds the custom all-gather size limit and falls back to a lazily initialized RCCL communicator, which showed up as a one-off multi-second TTFT stall on the first large CP+BCG prefill.
… for DSV4 on ROCm sgl-project#44022 dropped (interleave, dsv4) from the supported CP BCG combinations because the capture path did not shard input_ids with the hidden states for DSV4 hash routing. This branch fills input_ids_global from the CP BCG input buffers, so add the combination back on ROCm, where it is validated. CUDA keeps it disabled until the same path is validated on B200. Restore test_dsv4_prefill_cp_bcg.py with the DSV4 CP BCG checks and make the supported-configuration check follow the platform.
AMD-yanfeiwang
changed the title
AMD-yanfeiwang
force-pushed
the
amd/dsv4-rocm-prefill-cp-bcg
branch
from
October 14, 2026 13:02
8bdae57 to
0ccacc8
Compare
AMD-yanfeiwang
added a commit
to AMD-yanfeiwang/sglang
that referenced
this pull request
Oct 15, 2026
…oject#43465 by At1a8) Applied sgl-project#43465 (head 5c5862f, author @At1a8, open) on top of the HIP CP+BCG port (sgl-project#43819). Conflict resolution in deepseek_v4_backend_hip_radix.py: both PRs pad swa_replay_start and pass cp_metadata.local_index to apply_cp_reindex; kept the single copy from sgl-project#43819, which also keeps num_write_tokens for breakable prefill graphs.
AMD-yanfeiwang
added a commit
to AMD-yanfeiwang/sglang
that referenced
this pull request
Oct 15, 2026
…raphs on ROCm Stacked on the HIP V4.1 prefill CP (sgl-project#43465) and the HIP CP+BCG port (sgl-project#43819): - HIP low-ratio BCG capture accepts attention CP (the late-layer tail of decoder SWA bounded replay is already rejected with prefill graphs by the hook). - A CP replay builds the low-ratio metadata from the static batch, whose out_cache_loc is padded to the bucket. The write fields (raw/c1/c2 out_loc) are built from the live prefix, since the eager low-ratio sources write only the live rows; init_forward_metadata_in_graph still pins the bucket-sized SWA store target read by the captured segments. - _forward_low_ratio_sources_cp reads the live logical rows; CP BCG sets total_seq_lens to the capture bucket. - Engram hashing in the CP branch runs at an eager break under BCG. - On HIP, interleave CP BCG capture buckets start at cp_size * 16, the HIP minimum CP tokens per rank, so every captured bucket shards.
AMD-yanfeiwang
added a commit
to AMD-yanfeiwang/sglang
that referenced
this pull request
Oct 15, 2026
…oject#43465 by At1a8) Applied sgl-project#43465 (head 5c5862f, author At1a8, open) on top of the HIP CP+BCG port (sgl-project#43819), plus its later test update (07efa4b) that drops the stale V4.1 HIP CP rejection case. Conflict resolution in deepseek_v4_backend_hip_radix.py: both PRs pad swa_replay_start and pass cp_metadata.local_index to apply_cp_reindex; kept the single copy from sgl-project#43819, which also keeps num_write_tokens for breakable prefill graphs.
AMD-yanfeiwang
added a commit
to AMD-yanfeiwang/sglang
that referenced
this pull request
Oct 15, 2026
…raphs on ROCm Stacked on the HIP V4.1 prefill CP (sgl-project#43465) and the HIP CP+BCG port (sgl-project#43819): - HIP low-ratio BCG capture accepts attention CP (the late-layer tail of decoder SWA bounded replay is already rejected with prefill graphs by the hook). - A CP replay builds the low-ratio metadata from the static batch, whose out_cache_loc is padded to the bucket. The write fields (raw/c1/c2 out_loc) are built from the live prefix, since the eager low-ratio sources write only the live rows; init_forward_metadata_in_graph still pins the bucket-sized SWA store target read by the captured segments. - _forward_low_ratio_sources_cp reads the live logical rows; CP BCG sets total_seq_lens to the capture bucket. - Engram hashing in the CP branch runs at an eager break under BCG. - On HIP, interleave CP BCG capture buckets start at cp_size * 16, the HIP minimum CP tokens per rank, so every captured bucket shards. - CP decode attention TP also slices the gfx950 MXFP8 weight copies (weight_scale_mx_e8m0, weight_scale_inv_mx, weight_bf16) that V4.1's 32x32 ue8m0 fp8 blocks are served from; slicing only weight and the block scales failed the first decode-graph capture of wq_b with mismatched scale shapes.
This was referenced Oct 18, 2026
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
#38603 added interleave prefill context parallelism (CP) with breakable CUDA graphs (BCG) for DeepSeek V4, but only in the CUDA backend. On ROCm, prefill graph capture fails under CP (
deepseek_v4_backend_hip_radix.py: assert cp_metadata is not None), so deployments must choose between CP and BCG; the MI355X InferenceX DSv4 recipe, for example, uses BCG at conc 1/4 and CP at conc 16.#44022 later disabled DSv4 CP + BCG because the capture path did not shard token IDs consistently with hidden states for hash routing. This PR fixes that in the shared CP BCG path and enables DSv4 CP + BCG on ROCm, for DeepSeek-V4-Pro and DeepSeek-V4-Flash. DeepSeek-V4.1-Flash is in #44013.
Modifications
deepseek_v4_backend_hip_radix.py): under prefill graphs, use the bucket-sizedout_cache_locfor KV writes, passcp_metadata.local_indextoapply_cp_reindex, and padswa_replay_startto the CP layout, as [BCG][DSV4-Flash][CP] Support DeepSeek V4 prefill with interleave CP and breakable CUDA graphs #38603 does for CUDA.layers/cp/bcg.py,prefill_cuda_graph_runner.py): the captured body gets CP-local positions andinput_ids, whileforward_batch.positionsstays global for the SWA ring store.input_ids_globalis filled in eager CP's layout so hash-routed MoE layers route the right tokens (the [CP] Automatically enable BCG for supported prefill CP configurations #44022 issue; follows [Feat] [NVIDIA] deepseek-v4.1: Support prefill CP with BCG in PD serving (2/N) #40333).layers/cp/utils.py,bcg.py,eager_runner.py): skip all-gathering the pre-mHC-head states unless they are read. This removes ~4x extra gather traffic and an ~8.5 s TTFT stall on the first large prefill.layers/cp/bcg.py): add(INTERLEAVE, DSV4)on ROCm. CUDA is unchanged.test_dsv4_prefill_cp_bcg.pyand updateamd/test_dsv4_hip_bcg_metadata.py.Usage
On ROCm, BCG is selected automatically for this configuration (log:
Enabling breakable prefill CUDA graph for supported prefill CP). Pass--cuda-graph-backend-prefill disabledto run eager CP instead.Test setup
lmsysorg/sglang-rocm:v0.5.21-rocm724-mi35x-20261011. Stock is main3b72b70ef4and this PR is0ccacc8055, both overlaid on the same image.benchmarks/single_node/srt-slurm-recipes/dsv4/sglang/mi355x-fp4-mtp/agentic.yaml):SGLANG_HACK_FLASHMLA_BACKEND=unified_kv_triton SGLANG_DSV4_UNIFIED_KV_FP8=1 --kv-cache-dtype fp8_e4m3, two-pool FP8 nope + bf16 rope on gfx950). Under CP the packed FP8 KV is all-gathered before the store (_materialize_cp_unified_fp8_kv), and prefill attention runsprefill_fp8_2buff.--chunked-prefill-size 16384,--mem-fraction-static 0.86.--max-running-requestsand--cuda-graph-max-bs-decodeare 2x conc, and the server is restarted per conc.SGLANG_DSV4_GROUPED_ASM=1at conc 16 for all arms.--cuda-graph-backend-prefill breakable--enable-prefill-cp --cp-strategy interleave --enable-cp-decode-attn-tp--cuda-graph-backend-prefill breakableinfx.bench_serving.benchmark_serving): random dataset,--random-range-ratio 0.8, num-prompts 10x conc, 2x conc warmups, request rate inf,--ignore-eos.Accuracy Tests
GSM8K: InferenceX eval (lm-eval gsm8k, 5-shot, 1319 questions) on the conc-16 server.
All three are within one stderr (±0.005). Greedy outputs of CP and CP+BCG agree on prompts from 2 to 30k tokens; the only divergences are at near-tie tokens.
Speed Tests (DeepSeek-V4-Pro, TP8)
TTFT is mean / P50 / P90 in ms. TPOT is the mean in ms. Throughput is output tok/s.
ISL 1024 / OSL 1024
ISL 8192 / OSL 1024
ISL 8192 / OSL 2 (prefill-bound)
Summary
AgentX (DeepSeek-V4-Pro-0813, TP8, DSpark gamma 6)
Setup
srt_agentic.sh, tracessemianalysis_cc_traces_weka_062126, 900 s per point, all pointssubmission_valid), using the recipe's own conc 1/4/16 points.SGLANG_ENABLE_WAR_BARRIER=1. DSpark with BCG needs the WAR barrier for correctness on ROCm; the recipe's BCG arms run without it.Results (TTFT and ITL: mean / P50 / P90 ms; throughput: tok/s per GPU)
Takeaways
DeepSeek-V4-Flash, TP4
Setup
ISL 8192 / OSL 1024
ISL 1024 / OSL 1024
ISL 8192 / OSL 2 (prefill-bound)
GSM8K (InferenceX eval): BCG 0.9689, CP 0.9659, CP+BCG 0.9697.
Takeaways
BF16 unified KV (DeepSeek-V4-Flash, TP4)
All data above uses FP8 unified KV. This PR's position and KV-gather changes also apply to the BF16 unified layout, so it was checked separately. Settings are the same as the V4-Flash runs, except
SGLANG_DSV4_UNIFIED_KV_FP8=0(server log:unified=True unified_fp8=False). Arms: CP and CP+BCG at conc 4 and 16, each run twice on fresh servers.DeepSeek-V4.1-Flash results are in #44013.
Checklist
test_dsv4_prefill_cp_bcg.py,test_cp_strategy_unit.py,test_prefill_cuda_graph_runner.py,test_eager_on_graph.pyandamd/test_dsv4_hip_bcg_metadata.py(82 passed after the rebase).CI States
Latest PR Test (Base): ❌ Run #38568557194
Latest PR Test (Extra): ❌ Run #38568556932
Latest PR Test (AMD ROCm 10): ❌ Run #38568557126