[WS2] Add CP-aware Attention contract and dispatch metadata - #236
[WS2] Add CP-aware Attention contract and dispatch metadata #236Bsgg1 wants to merge 12 commits into
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds a typed WS2 CP-aware attention contract, validates attention metadata, adds capability-aware registry dispatch with provenance, and documents and tests strict backend rejection behavior. ChangesWS2 attention contract and dispatch
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/attention_contract.py`:
- Around line 403-420: Rename the loop-bound active_blocks variables in the
sequence consistency loop and the exclusive_page_owners loop to distinct names,
updating their references within each loop. Preserve the existing list-based
active_blocks binding and all validation and ownership logic.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: bb3bdb48-0e1e-4f2b-9707-bebdb8e47072
📒 Files selected for processing (6)
docs/design/runtime-dispatch.mddocs/design/ws2-cp-attention-contract.mddocs/operators/attention.mdrl_engine/kernels/attention_contract.pyrl_engine/kernels/registry.pytests/test_attention_contract.py
| object.__setattr__(self, "mode", _enum_value(AttentionMode, self.mode, "mode")) | ||
| object.__setattr__(self, "dtype", _enum_value(AttentionDType, self.dtype, "dtype")) | ||
| batch_size = _positive_int(self.batch_size, "batch_size") | ||
| _positive_int(self.query_sequence_length, "query_sequence_length") |
There was a problem hiding this comment.
Small contract gap: for prefill, should we reject cases where query_sequence_length doesn’t match sharding.local_sequence_length? Right now an inconsistent contract like query_sequence_length=2048 with local_sequence_length=1024 can still be constructed.
| global_token_positions = _integer_tuple( | ||
| self.global_token_positions, "global_token_positions" | ||
| ) | ||
| if not cache_positions or any(position < 0 for position in cache_positions): |
There was a problem hiding this comment.
Another small gap: cache_positions is only checked for shape/non-negativity, but not against global_token_positions. For decode identity, we probably want to reject cases like cache_positions=(999,) when the declared tokens are only 0..16.
There was a problem hiding this comment.
🧹 Nitpick comments (2)
tests/test_attention_contract.py (2)
458-466: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExercise registry rejection for packed-layout incompatibility.
This test only calls
AttentionBackendCapability.incompatibilities(...); it never proves thatKernelRegistry.get_attention_op()rejects an otherwise declared backend. Register the capability and assert the dispatch error includes the packed-varlen incompatibility, preventing registry regressions from bypassing this check.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_attention_contract.py` around lines 458 - 466, Extend test_packed_layout_requires_declared_backend_support to register the capability with KernelRegistry and invoke get_attention_op() for the packed-varlen contract. Assert that dispatch is rejected and the resulting error includes “packed varlen layout is unsupported,” while retaining the direct incompatibilities assertion.
100-101: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCover the stated TP=4/CP=4 target, not only TP=2/CP=2.
The representability test and its declared backend only exercise TP2/CP2, while the PR objective explicitly calls out a Qwen3-8B TP4/CP4 BF16 contract. Add a TP4/CP4 case with matching local sequence/block metadata and capability support so four-way head ownership and CP validation are covered.
Also applies to: 110-116
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/test_attention_contract.py` around lines 100 - 101, Extend the representability test’s world-size parameters to include the stated TP=4/CP=4 scenario, and add matching local sequence/block metadata plus backend capability support for that case. Update the relevant test setup around the representability test and its declared backend so four-way tensor-parallel head ownership and context-parallel validation are exercised while preserving the existing TP2/CP2 coverage.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@tests/test_attention_contract.py`:
- Around line 458-466: Extend
test_packed_layout_requires_declared_backend_support to register the capability
with KernelRegistry and invoke get_attention_op() for the packed-varlen
contract. Assert that dispatch is rejected and the resulting error includes
“packed varlen layout is unsupported,” while retaining the direct
incompatibilities assertion.
- Around line 100-101: Extend the representability test’s world-size parameters
to include the stated TP=4/CP=4 scenario, and add matching local sequence/block
metadata plus backend capability support for that case. Update the relevant test
setup around the representability test and its declared backend so four-way
tensor-parallel head ownership and context-parallel validation are exercised
while preserving the existing TP2/CP2 coverage.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 19a56bfb-94c2-40c6-b17b-1422b2d0c29f
📒 Files selected for processing (2)
docs/design/ws2-cp-attention-contract.mdtests/test_attention_contract.py
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/design/ws2-cp-attention-contract.md
Signed-off-by: inaniloquentee <3051000145@qq.com>
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/attention_contract.py`:
- Around line 606-618: The RoPE validation around position_ids and
query_position_offsets/key_position_offsets currently checks only lengths and
must enforce CP token ownership. Require position metadata for RoPE-aware
requests, validate dense position_ids against global_block_token_starts and
local_block_offsets, and reject offset-only metadata when ownership is
non-contiguous; add regression coverage for blocks with non-contiguous global
starts such as (0, 3072).
- Around line 495-500: Update theta validation in the attention contract around
self.theta to reject non-finite values after converting to float. Require
math.isfinite(theta) alongside the existing positive-value check before
object.__setattr__ stores the normalized value, while preserving the current
error behavior and message.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: b24f9885-0127-42d0-bd01-c402cfe94026
📒 Files selected for processing (3)
docs/design/ws2-cp-attention-contract.mdrl_engine/kernels/attention_contract.pytests/test_attention_contract.py
| if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): | ||
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | ||
| theta = float(self.theta) | ||
| if theta <= 0.0: | ||
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | ||
| object.__setattr__(self, "theta", theta) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Reject non-finite theta values.
float("nan") and float("inf") pass theta <= 0.0. They can produce invalid RoPE semantics and non-standard JSON values in provenance. Require math.isfinite(theta) before storing the value.
Proposed fix
- if theta <= 0.0:
+ if not math.isfinite(theta) or theta <= 0.0:
raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}")📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): | |
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | |
| theta = float(self.theta) | |
| if theta <= 0.0: | |
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | |
| object.__setattr__(self, "theta", theta) | |
| if isinstance(self.theta, bool) or not isinstance(self.theta, (float, int)): | |
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | |
| theta = float(self.theta) | |
| if not math.isfinite(theta) or theta <= 0.0: | |
| raise AttentionContractError(f"theta must be a positive number; got {self.theta!r}") | |
| object.__setattr__(self, "theta", theta) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/attention_contract.py` around lines 495 - 500, Update theta
validation in the attention contract around self.theta to reject non-finite
values after converting to float. Require math.isfinite(theta) alongside the
existing positive-value check before object.__setattr__ stores the normalized
value, while preserving the current error behavior and message.
| if self.rope.position_ids is not None and len(self.rope.position_ids) not in { | ||
| query_sequence_length, | ||
| self.sharding.local_sequence_length, | ||
| }: | ||
| raise AttentionContractError( | ||
| "position_ids must describe the local query sequence or full local " | ||
| "sequence length" | ||
| ) | ||
| for field in ("query_position_offsets", "key_position_offsets"): | ||
| offsets = getattr(self.rope, field) | ||
| if offsets is not None and len(offsets) != batch_size: | ||
| raise AttentionContractError( | ||
| f"{field} must contain one entry per logical batch entry" |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
Bind RoPE positions to CP token ownership.
The current checks only validate lengths. They accept absent position metadata and accept position_ids=range(local_sequence_length) for non-contiguous CP blocks. For example, blocks with global starts (0, 3072) would assign the second local block positions starting at 1024, not 3072. This changes the RoPE phase and can produce incorrect attention results.
Require position identity for RoPE-aware requests. Validate dense IDs against global_block_token_starts and local_block_offsets. Reject an offset-only representation when it cannot represent non-contiguous ownership. Add regression coverage for this case.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@rl_engine/kernels/attention_contract.py` around lines 606 - 618, The RoPE
validation around position_ids and query_position_offsets/key_position_offsets
currently checks only lengths and must enforce CP token ownership. Require
position metadata for RoPE-aware requests, validate dense position_ids against
global_block_token_starts and local_block_offsets, and reject offset-only
metadata when ownership is non-contiguous; add regression coverage for blocks
with non-contiguous global starts such as (0, 3072).
# Conflicts: # docs/operators/attention.md # rl_engine/kernels/gtest/operator_specs.py # rl_engine/kernels/ops/cuda/attention/__init__.py # rl_engine/kernels/registry.py
Signed-off-by: lamentropetion <3051000145@qq.com>
Signed-off-by: lamentropetion <3051000145@qq.com>
2 participants
Summary
Implements PR1 of #235: the typed WS2 Attention contract and strict dispatch metadata used by the later CP reference, runtime binding, decode, backward, and H100 validation PRs.
Contract
(Out, LSE)merge and final-write downcastbatch x TP x CP x ownerplan provenance; missing/mismatched evidence fails closedrlkernel.attention.deterministic_core.v1as the strict production core identityValidation
45 passed67 passedtestBoundary
Multi-rank execution evidence is owned by #279 and accepted by #285. Related: #235, #108, #207.