[WS2] Add CP-aware Attention contract and dispatch metadata - #236
Conversation
📝 WalkthroughWalkthroughThis PR adds a typed WS2 CP-aware attention contract, validates sharding, reduction, causal, packed-varlen, and KV-cache 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 issues
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
2 participants
Summary
Implements PR 1 of #235 by introducing the explicit WS2 contract and dispatch metadata for CP-aware deterministic standard softmax Attention.
This PR defines how TP/CP sharding, deterministic
(out, lse)reduction, packed sequences, and decode KV-cache identity must be described andvalidated before a backend can be selected.
Changes
AttentionContractShardingSpecReductionSpecKVCacheSpecAttentionBackendCapabilityKernelRegistry.get_attention_op()dispatch.Behavior
A valid Qwen3-8B
TP=4,CP=4, BF16 contract can now be represented and validated.Existing WS1 Attention backends are intentionally rejected for strict WS2 requests because they do not yet export attention-domain LSE or implement
deterministic CP
(out, lse)merge.Legacy
KernelRegistry.get_op()behavior is unchanged.Validation