Zigzag cp - #3
Open
alexsun07 wants to merge 4 commits into
Open
Conversation
Wires the AITER attention backend into SGLang's existing prefill CP infrastructure (cp_allgather_and_save_kv_cache + cp_attn_forward_extend from cp_utils.py), enabling --enable-prefill-context-parallel to work with --attention-backend aiter on ROCm. Three hunks in aiter_backend.py: 1. Import the CP utilities. 2. Cache attn_cp_size on AiterAttnBackend. 3. In forward_extend: detect CP mode, route KV writes through the CP allgather wrapper (so each rank's local pool holds the full sequence K/V), and split prefill attention into prev/next zigzag halves via cp_attn_forward_extend, calling mha_batch_prefill_func per half. Verified on Qwen/Qwen3-30B-A3B-FP8 with these working configs (all require moe_tp_size == 1 to avoid an upstream MoE-TP/CP interaction unrelated to attention): tp=2 cp=2 moe_dp=2 tp=4 cp=4 moe_dp=4 tp=4 cp=2 moe_dp=2 ep=2 tp=4 cp=2 moe_dp=1 ep=4 (matches TestQwen330BCP) tp=8 cp=2 moe_dp=2 ep=4 Co-Authored-By: Claude Opus 4 (1M context) <noreply@anthropic.com>
Three small follow-ups on top of the initial AITER prefill CP wiring: 1. Match FA's init style for attn_cp_size (drop the defensive getattr; model_runner.attn_cp_size is unconditionally set). 2. Pre-allocate a 2-element kv_indptr scratch tensor on the backend when attn_cp_size > 1 and update it in place per zigzag half. Avoids building a fresh torch.tensor for every mha_batch_prefill_func call (~2 per layer × 2 halves per Qwen3-30B forward). 3. Remove the int(cache_seqlens.item()) host sync from the CP closure. AITER's mha_batch_prefill_func requires max_seqlen_k as a Python int while FA's flash_attn_with_kvcache accepts a tensor, so the shared cp_attn_forward_extend helper only hands backends the tensor form. Inline the prev/next zigzag split directly in aiter_backend so we can source kv_len_prev / kv_len_next from cp_meta as Python ints (already populated in prepare_context_parallel_metadata). Same call structure as cp_attn_forward_extend, just bypasses the abstraction. Drops the unused cp_attn_forward_extend import. Verified bit-identical generation vs the previous commit (smoke prompt) and no GSM8K accuracy regression on all four working CP configs: tp=2/cp=2, tp=4/cp=4, tp=4/cp=2 ep=2 (path A), tp=4/cp=2 moe_dp=1 ep=4 (path B; matches upstream TestQwen330BCP). Cuda graph remains disabled in launch recipes — a pre-existing AITER decode + cuda graph correctness issue (reproduces without CP) is being tracked separately. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Registers test/registered/4-gpu-models/test_qwen3_30b.py for the stage-c-test-4-gpu-amd suite alongside the existing CUDA suite. On AMD the default attention backend is AITER, so TestQwen330BCP exercises the AITER prefill CP path (the same primitives FA exercises on CUDA: KV allgather + zigzag prev/next halves) without changing how the test is written. Adds one platform-conditional argument: --disable-cuda-graph on HIP. This is needed because of a separate AITER + cuda graph correctness issue on Qwen3-MoE decode that reproduces independently of CP. The CUDA path is unchanged byte-for-byte (the helper returns an empty list when not is_hip()). Verified locally on MI300X: TestQwen330BCP::test_gsm8k passes (GSM8K 200 examples, gate score >= 0.85), wall time ~15 min. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Replaces the conservative --disable-cuda-graph workaround with the
narrower fix found by separate AITER investigation:
1. --disable-custom-all-reduce: AITER's custom all-reduce has a
relaxed-memory-ordering data race that fires when the entire decode
forward is captured into a single cuda graph (allreduce reads can
land before peer writes). Falling back to NCCL for the allreduce
sidesteps it; cuda graph itself works fine.
2. TORCH_NCCL_BLOCKING_WAIT=1: when the attention CP process group
exists, PyTorch's NCCL watchdog races with cuda-graph capture on
HIP — it polls hipEventQuery on Work whose completion event lives
on a stream currently in capture mode, which HIP forbids
(`hipErrorCapturedEvent`). Switching to blocking-wait mode
disables the watchdog thread and lets capture complete. Not
needed on CUDA — the FA path's CP test runs without it.
End-to-end GSM8K(200) on MI300X: 3m38s with cuda graph vs 14m38s
without. Passes the GSM8K_BASELINE_ACCURACY gate (>=0.85).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
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
1 participant
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
Modifications
Accuracy Tests
Benchmarking and Profiling
Checklist
Review Process
/tag-run-ci-label,/rerun-failed-ci,/tag-and-rerun-ci