[WS2] Add CP-aware Attention contract and dispatch metadata - #236
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 Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
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>
…PR4) Wire the CP attention path into the cross-configuration planner/runtime for the Qwen3-8B TP=2 CP=2 BF16 target. The PR4 criterion "rollout and training descriptors bind to the same semantic attention contract" cannot hold literally: training runs full-sequence prefill over a CP-sharded sequence while rollout runs vLLM paged-KV chunked prefill, so the two AttentionContract instances always differ. Binding is therefore split into three tiers -- identity must match bit for bit, reduction semantics must match each other and the WS2 mandate, and materialization differences are recorded and measured rather than rejected. reduction.engine stays in the recorded tier so a Transformer Engine merge oracle on one side does not fail the binding; reduction.order and acc_dtype stay in the semantic tier because that is the WS2 claim. Also adds the first two framework-shaped RuntimeMaterializer implementations. Before this the only one was CpuSmokeMaterializer over a synthetic CPU model, and every named scenario was planning-only. Neither adapter imports megatron or vllm, so the binding rules run on CPU in CI. Determinism is probed on both sides and compared, because the two frameworks mean different things by it: Megatron asserts NCCL_ALGO and leaves TF32 and BF16 reduced-precision reduction unmanaged, while vLLM hard-sets ten NCCL variables and disables both. Mismatches in NCCL_ALGO, NCCL_PROTO and CUBLAS_WORKSPACE_CONFIG are blocking; the rest are recorded. Fixes a latent break on the way: the planner normalizes dtype knobs to torch spellings (bfloat16) while AttentionDType uses short ones (bf16), so passing a normalized knob into the enum raised. Stacked on #236 (attention contract) and #238 (deterministic CP reference), on top of #230 (cross-configuration framework). Part of #235 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw
Three fields could differ between the two sides without the binding noticing. dtype was in no tier at all, so a BF16 rollout could bind to an FP16 training pass and produce a drift number attributable to nothing. It joins the semantic tier, with allow_dtype_difference for the #235 PR5 sweep that deliberately scores BF16 against an FP32 reference. batch_size was likewise unchecked. Batch invariance is a claim about results not changing with batch makeup, so two sides scoring different batches are not comparable and it belongs to identity. split_kv_policy has no field in the #236 contract, so it only reached side_configs and never took part in binding. Callers now pass it through rollout_recorded_extra / training_recorded_extra so the difference is at least visible in provenance; it can move into the contract once #236 grows the field. Part of #235 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw
3 participants
Scope
PR1 defines the typed Attention contract and strict dispatch metadata used by the later CP reference, decode, backward, and acceptance PRs.
Contract
Validation
Boundary
This PR owns contract construction and capability gating. CP execution, deterministic arithmetic, and H100 evidence are validated by PR279 and PR285.
Commit: 51b8f52