feat(attention): add CP backward drift validation - #284
Merged
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
inaniloquentee
changed the base branch from
feat/ws2-cp-attention-reference-pr3
to
test
August 12, 2026 16:13
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com> # Conflicts: # tests/test_operator_inputs.py
maxiaosong1124
approved these changes
Aug 24, 2026
zhangj1an
added a commit
that referenced
this pull request
Aug 24, 2026
Also makes the strict path TP-degree bitwise invariant. Merge resolution (12 conflicted files): - registry.py: test grew its own get_attention_op / _attention_policy_mismatch / _platform / _get_or_create_backend, keyed off a `ws2_attention` priority list and a flat _attention_capabilities dict. Dropped our parallel implementation and rebuilt on test's, keeping only the register_attention_ backend seam (a vendor-conditional backend cannot be declared in a static list) and the auto-under-CP guard. Corrects an error on our side: requested_backend="deterministic" was rejected by copying a logprob-dispatch rule, but AttentionBackendCapability admits "deterministic" and it is test's default. - csrc/ops.cpp: test registers the CUDA-IPC collectives unguarded, but their declarations stay ROCm-guarded, so taking that side fails to compile on ROCm. Kept our guard, which is the version this ROCm build actually succeeded with. - flashinfer_paged_attention.py (9 hunks), ws2_p2p_nccl_*.py (3): test reverted to CUDA-only expressions (hardcoded cuda_ag_rs, direct StrictFlashAttention4 Core). Ours is the platform-aware superset. - cp_attention.py: took test's new saved-forward-state validation (#284). - test_flashinfer_pr7_attention.py: the two sides were different tests at the same position (ROCm core accepted vs reference core rejected). Kept both. - setup.py and pyproject.toml are EMPTY on test (0 bytes, from 3e04a63) and pyproject.toml auto-merged to empty with no conflict. Restored both; test cannot build a native extension in its current state. TP-degree invariance: every launch is now pinned to one batch row and one KV group. AITER's reduction order is launch-shape dependent, so a head shard computed under TP=4 was not bit-identical to the same shard under TP=8 (up to 7.8125e-03 at some shapes). Pinning removes it: out/lse/dQ/dK/dV are bitwise equal at TP=2/4/8 across S=512..4096, at roughly 3x forward time. The cross-config contract therefore no longer binds the TP degree -- doing so would reject comparisons that are in fact identical -- and checks only what changes the arithmetic. ROCm test skips: CUDA reports through the same device API on ROCm, so device_count guards do not exclude it and CUDA-exclusive tests failed instead of skipping. Added a `cuda_only` marker plus a conftest hook, applied to the CUDA-IPC collectives, the CUDA det_gemm K-tree cases, and the FA4-selection test (ROCm correctly resolves to the AITER core there, so its monkeypatch is never consulted). Attention suites: 323 passed, 42 skipped. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfN3b2ep5DvxVDir7J36j3
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
3 participants
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.
Scope
PR8 fixes the strict backward schedule for CP prefill and chunked-prefill.
Contract
Validation
dQ/dK/dVmax abs0.0across 2/4/8 ranksBoundary
This PR owns backward arithmetic order and gradient coverage. PR285 and #294 record the production H100 evidence.
Commit: 9e5c400