Repository navigation
[AMD][DSV4.1][CP] Support interleave prefill CP with breakable CUDA graphs on ROCm - #44013
Open
AMD-yanfeiwang wants to merge 5 commits into
Open
AMD-yanfeiwang wants to merge 5 commits into
AMD-yanfeiwang wants to merge 5 commits into
Conversation
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
3 tasks done
…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
force-pushed
the
amd/dsv41-rocm-prefill-cp-bcg
branch
from
October 15, 2026 11:53
0c983ac to
beb70b2
Compare
…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.
…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.
AMD-yanfeiwang
force-pushed
the
amd/dsv41-rocm-prefill-cp-bcg
branch
from
October 15, 2026 12:00
beb70b2 to
9cba916
Compare
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
This extends ROCm interleave prefill CP + breakable CUDA graphs (BCG) to DeepSeek-V4.1-Flash. On HIP, V4.1 uses the low-ratio (C1/C2) layout, and its breakable prefill graph path rejects attention CP.
Stacked PR. Please review only the last commit.
5d062ce956and will be dropped once [Feat] [AMD] DeepSeek-v4.1: Support prefill context parallel #43465 merges.7d7e991957.Modifications (last commit)
deepseek_v4_backend_hip_radix.pyout_cache_loc.deepseek_v4.py_forward_low_ratio_sources_cpuses the live logical token count. Under CP BCG,total_seq_lensis the capture bucket.layers/cp/bcg.py: on HIP, CP BCG capture buckets start atcp_size * 16, the HIP minimum CP tokens per rank from [Feat] [AMD] DeepSeek-v4.1: Support prefill context parallel #43465. This ensures every captured bucket is CP-active.test_dsv4_prefill_cp_bcg.pyandtest_cp_strategy_unit.py(38 passed; 14 of the new tests fail without this commit).Accuracy Tests
Setup:
lmsysorg/sglang-rocm:v0.5.21-rocm724-mi35x-20261011Arms:
--cuda-graph-backend-prefill breakable--enable-prefill-cp --cp-strategy interleave --cuda-graph-backend-prefill disabled--cuda-graph-backend-prefill breakableThe CP arms need
--language-model-only, because V4.1 vision does not support CP. The stock image rejects this flag, so the BCG arm runs without it. A control run of this branch, BCG with the flag, matches stock BCG within 1%.Speed Tests (preliminary)
Prefill TTFT p50 (ms), ISL as listed, OSL 1, measured with
bench_servingrandom. This is preliminary data: InferenceX-recipe and client results (fixed-seq 1k/1k, 8k/1k) will be added.Serving 8k/1k: TTFT, TPOT and throughput are within ±5% across arms. Decode already uses graphs in all of them.
Takeaways:
Checklist
CI States
Latest PR Test (Base): ❌ Run #38731361472
Latest PR Test (Extra): ❌ Run #38731361598
Latest PR Test (AMD ROCm 10): ❌ Run #38731361474