docs(ws1): batch-invariance contract RFC (#101) - #264
Conversation
Defines the WS1 batch-invariance contract: the same logical sequence must produce identical active-token outputs regardless of surrounding batch shape. Separates three comparison relations so the threshold for any check follows from its kind rather than from negotiation: I (invariance) kernel vs itself under a shape transform -> bitwise A (accuracy) kernel vs fp32 gold -> RL-Align#108 table P (parity) backend vs backend -> RL-Align#108 table Bitwise is required on I because only the batch shape moves: the kernel, dtype, data and device are held fixed, so no legitimate source of difference exists and a tolerance would hide the defect the suite exists to find. This policy is already declared as batch_invariance {atol: 0, rtol: 0} in tolerance_contract.json but is read by no code path; _resolve_tolerance only consults contract["accuracy"]. Names six transform classes (T1 population, T2 position, T3 neighbour content, T4 padding, T5 execution schedule, T6 selection/packing) as source_class values compatible with the WS2 KnobDefinition model, so a WS1 axis lifts into a WS2 ablation without redefinition. Positions WS1 as the precondition for a trustworthy WS2 A0 reference. Adds a drift report schema, chain-index fault localisation, the P0 suite matrix with P1/P2 deferrals, and local plus CI command shapes. Refs RL-Align#101, RL-Align#83, RL-Align#108, RL-Align#96, RL-Align#111 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
📝 WalkthroughWalkthroughAdded the WS1 batch-invariance RFC. It defines invariance transforms, comparison rules, fixtures, reporting, harness integration, CI commands, failure classification, follow-up issues, and completion criteria. ChangesBatch-invariance contract
Estimated code review effort: 1 (Trivial) | ~5 minutes Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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 |
There was a problem hiding this comment.
Pull request overview
Adds an RFC design document defining WS1’s batch-invariance contract for RL-Kernel’s train–inference consistency suite, clarifying that invariance checks (shape-only transforms) must be bitwise while accuracy/parity checks remain tolerance-based per the WS1 numerical contract (#108).
Changes:
- Defines the WS1 invariance (I) / accuracy (A) / parity (P) relations and their thresholds, emphasizing bitwise equality for invariance.
- Specifies transform taxonomy (T1–T7), the P0 suite matrix, drift report schema, and first-divergence localization strategy.
- Documents intended harness integration (
run_invariance_suite) and contributor/CI command shapes plus a follow-up issue breakdown.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| "mean_abs_error": 4.1e-07, | ||
| "n_mismatched": 12, | ||
| "n_active": 4096, | ||
| "first_mismatch": {"sequence_row": 3, "token_position": 117, "vocab_index": 40213}, |
| OperatorSpec.invariance: InvarianceDeclaration | None | ||
| chain_index: int | ||
| transforms: frozenset[str] # which classes apply | ||
| backward: bool | ||
| exempt_reason: str | None # required when invariance is None |
| op_class: elementwise | reduction | attention | logprob | ||
| dtype: float32 | bfloat16 | float16 | ||
| device: cpu | cuda | rocm | ||
| seed: int |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 10
🤖 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 `@docs/design/ws1_batch_invariance_contract.md`:
- Line 300: Fix the Markdown heading syntax in the WS1 batch invariance contract
by adding a space after the hash marker so the heading renders correctly and
satisfies MD018.
- Around line 537-539: Remove the unused Markdown reference definitions for `#83`
and `#101` from the document, while preserving the existing inline URLs and
leaving the still-referenced `#96` definition unchanged.
- Around line 431-444: Move exempt_reason into the OperatorSpec state, or model
OperatorSpec as a discriminated union containing either an invariance
declaration or an exemption reason. Ensure operators with invariance=None can
carry and require the exemption reason, while covered operators retain their
existing invariance fields.
- Around line 260-296: Update the Drift Report schema and accompanying
description to make first_mismatch operator-generic: replace the
vocabulary-specific physical-row fields with a stable logical sequence
identifier, explicit baseline and variant row positions, and an
operator-specific output coordinate or shape. Define whether n_mismatched and
n_active count elements, coordinates, tokens, or another unit, and ensure the
example reflects these stable semantics across T2 and non-vocabulary operators.
- Around line 304-317: Clarify the output contract for the loss-reduction
operators grpo_loss and ratio_kl in the invariance chain: define them as
per-token outputs or specify a fixed, batch-independent normalization before
requiring T1, T3, and T4 bitwise invariance. Apply the same clarification to the
corresponding section around the additional loss-reduction requirements.
- Around line 217-239: Extend InvarianceFixture.logical and the replay contract
with deterministic backward inputs, including the upstream gradient/backward
seed and any saved forward state required by the operator. Ensure backward
checks reuse these recorded values across baseline and transformed executions;
alternatively, explicitly restrict backward invariance checks to operators with
a fixed implicit gradient.
- Around line 136-143: Resolve the T5 scope inconsistency in the contract:
either define the schedule transformations and outputs that must remain
invariant for det_gemm, including the P0 matrix and follow-up requirements, or
explicitly mark T5 as not applicable to det_gemm. Update the related T5
references consistently so the attention-specific wording does not conflict with
the stated implementation scope.
- Around line 397-411: Update the RFC language to distinguish planned work from
existing functionality: change “Gaps this RFC closes” to describe gaps the RFC
identifies, and mark the generator, runner, report, and CLI command sections as
proposed or post-implementation. Apply the same terminology consistently in the
referenced follow-up table and command sections.
- Around line 319-332: Update the reporting contract for the failure records and
primary triage verdict so every independently failing operator remains visible
and classified, including higher-chain-index failures when an earlier operator
also fails. Keep the lowest chain_index as the primary investigation result, but
explicitly distinguish downstream failures caused by propagation from
independently observed operator failures; apply the same rule to the related
section around the failure-record requirements.
- Around line 217-254: Update the fixture_id and replay contract in the
InvarianceFixture specification so fixture_id alone can deterministically
reconstruct the complete fixture, including seed, generator version,
logical/tensor inputs, and every transform parameter; alternatively require a
serialized fixture manifest and make reports reference it. Align the replay and
serialization requirements, including the corresponding repeated section, so
storing seed separately is no longer presented as sufficient.
🪄 Autofix
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 Plus
Run ID: 0ad53b98-c719-4326-bbd8-b655bb944af6
📒 Files selected for processing (1)
docs/design/ws1_batch_invariance_contract.md
| ### T5: Execution schedule | ||
|
|
||
| Change how the sequence is split across launches without changing the sequence: chunked-prefill | ||
| boundary placement, prefill/decode split points, and KV-cache handoff. Scoring `[0:S]` in one | ||
| launch must equal scoring `[0:C]` then `[C:S]` against the resulting cache. | ||
|
|
||
| This class governs attention specifically, including the exported LSE, which must be invariant on | ||
| the same terms as the attention output. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Resolve the T5 scope mismatch.
T5 says it governs attention specifically, but the P0 matrix and follow-up issue require T5 for det_gemm. Define the det_gemm schedule transforms and compared outputs, or mark T5 as not applicable to det_gemm. The current text gives conflicting implementation scope.
Also applies to: 340-346, 507-515
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 136 - 143, Resolve
the T5 scope inconsistency in the contract: either define the schedule
transformations and outputs that must remain invariant for det_gemm, including
the P0 matrix and follow-up requirements, or explicitly mark T5 as not
applicable to det_gemm. Update the related T5 references consistently so the
attention-specific wording does not conflict with the stated implementation
scope.
| An invariance fixture is a logical sequence plus a declared set of transforms. The generator's job | ||
| is to guarantee that applying a transform changes only what the transform names. | ||
|
|
||
| ```text | ||
| InvarianceFixture: | ||
| fixture_id: content-derived stable id | ||
| op_name: key into OP_SPECS | ||
| op_class: elementwise | reduction | attention | logprob | ||
| dtype: float32 | bfloat16 | float16 | ||
| device: cpu | cuda | rocm | ||
| seed: int | ||
|
|
||
| logical: | ||
| token_ids: [S] or [B, S] | ||
| position_ids: optional, explicit when the op consumes them | ||
| target_ids: for logprob ops | ||
| attention_mask: [B, S] | ||
| action_mask: [B, S], defines the active token set | ||
| tensor_inputs: op-specific, generated from seed | ||
|
|
||
| reference_row: which row of the baseline holds the sequence under test | ||
|
|
||
| transforms: list of declared transform instances (T1..T6) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Specify deterministic inputs for backward checks.
The contract requires backward invariance, but InvarianceFixture.logical has no upstream gradient, backward seed, or saved-state definition. Backward outputs can change when grad_output changes even when forward inputs are identical. Add these values to the fixture and replay contract, or limit backward checks to operators with a fixed implicit gradient.
Also applies to: 340-352
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 217 - 239, Extend
InvarianceFixture.logical and the replay contract with deterministic backward
inputs, including the upstream gradient/backward seed and any saved forward
state required by the operator. Ensure backward checks reuse these recorded
values across baseline and transformed executions; alternatively, explicitly
restrict backward invariance checks to operators with a fixed implicit gradient.
| An invariance fixture is a logical sequence plus a declared set of transforms. The generator's job | ||
| is to guarantee that applying a transform changes only what the transform names. | ||
|
|
||
| ```text | ||
| InvarianceFixture: | ||
| fixture_id: content-derived stable id | ||
| op_name: key into OP_SPECS | ||
| op_class: elementwise | reduction | attention | logprob | ||
| dtype: float32 | bfloat16 | float16 | ||
| device: cpu | cuda | rocm | ||
| seed: int | ||
|
|
||
| logical: | ||
| token_ids: [S] or [B, S] | ||
| position_ids: optional, explicit when the op consumes them | ||
| target_ids: for logprob ops | ||
| attention_mask: [B, S] | ||
| action_mask: [B, S], defines the active token set | ||
| tensor_inputs: op-specific, generated from seed | ||
|
|
||
| reference_row: which row of the baseline holds the sequence under test | ||
|
|
||
| transforms: list of declared transform instances (T1..T6) | ||
| ``` | ||
|
|
||
| Requirements: | ||
|
|
||
| 1. **Transform isolation.** Applying a transform must change only its declared dimension. The | ||
| generator emits the baseline and the transformed batch from the same seed and the same logical | ||
| tensors; it must not regenerate random inputs per case. An undeclared change invalidates the | ||
| case, mirroring the WS2 `IsolationValidator` rule. | ||
| 2. **Extraction is exact.** Pulling the sequence under test out of a transformed batch is indexing, | ||
| never recomputation. No gather that reorders, no cast, no contiguity change that could alter the | ||
| compared values. | ||
| 3. **Active-token masking is the generator's responsibility.** Comparisons receive already-masked | ||
| tensors so that no downstream check accidentally compares padding. | ||
| 4. **Deterministic and serializable.** `fixture_id` is derived from content, so a failing case can | ||
| be re-run standalone from its id and reported in an issue without attaching tensors. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Make fixture_id sufficient for replay.
The RFC requires replay from fixture_id alone, but the example identifier only contains a short content-derived value. It does not encode or reference the seed, generator version, tensor inputs, or all transform parameters. The report stores seed separately, so the stated replay guarantee is not met. Define an invertible fixture specification, or require a serialized fixture manifest for replay.
Also applies to: 524-531
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 217 - 254, Update
the fixture_id and replay contract in the InvarianceFixture specification so
fixture_id alone can deterministically reconstruct the complete fixture,
including seed, generator version, logical/tensor inputs, and every transform
parameter; alternatively require a serialized fixture manifest and make reports
reference it. Align the replay and serialization requirements, including the
corresponding repeated section, so storing seed separately is no longer
presented as sufficient.
| ## Drift Report | ||
|
|
||
| Every I-relation failure emits one record. The schema is fixed so that CI, local runs, and issue | ||
| reports are the same artifact. | ||
|
|
||
| ```json | ||
| { | ||
| "relation": "invariance", | ||
| "status": "fail", | ||
| "fixture_id": "logp-bf16-cuda-a91c3f", | ||
| "operator": "batch_invariant_logp", | ||
| "op_class": "logprob", | ||
| "chain_index": 8, | ||
| "transform": {"class": "T1", "name": "batch_population", "baseline": 1, "variant": 16}, | ||
| "backend": "cuda", | ||
| "arch_key": "sm90", | ||
| "device": "cuda:0", | ||
| "dtype": "torch.bfloat16", | ||
| "direction": "forward", | ||
| "threshold": {"atol": 0.0, "rtol": 0.0}, | ||
| "max_abs_error": 3.0517578125e-05, | ||
| "mean_abs_error": 4.1e-07, | ||
| "n_mismatched": 12, | ||
| "n_active": 4096, | ||
| "first_mismatch": {"sequence_row": 3, "token_position": 117, "vocab_index": 40213}, | ||
| "launch": {"batch_size": 16, "seq_len": 256, "chunk_size": null, "padding_side": "right"}, | ||
| "seed": 123 | ||
| } | ||
| ``` | ||
|
|
||
| `first_mismatch` reports the lowest `(sequence_row, token_position)` in index order, not the | ||
| largest error. The first divergence is the debugging entry point; the largest error is usually | ||
| downstream of it. | ||
|
|
||
| The existing `OutputCheck` dataclass in `op_checks.py` already carries `max_abs_error`, | ||
| `mean_abs_error`, `max_rel_error`, and `passed`. This schema is that structure plus transform, | ||
| backend, launch, and position metadata. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Make drift coordinates stable and operator-generic.
first_mismatch assumes a vocabulary-like output and reports a physical row. T2 changes the physical row, and most P0 operators do not have a vocab_index. Define a stable logical sequence identifier, baseline and variant row positions, and an operator-specific output coordinate or shape. Also define the units for n_mismatched and n_active.
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 260 - 296, Update
the Drift Report schema and accompanying description to make first_mismatch
operator-generic: replace the vocabulary-specific physical-row fields with a
stable logical sequence identifier, explicit baseline and variant row positions,
and an operator-specific output coordinate or shape. Define whether n_mismatched
and n_active count elements, coordinates, tokens, or another unit, and ensure
the example reflects these stable semantics across T2 and non-vocabulary
operators.
|
|
||
| ## First-Divergence Localization | ||
|
|
||
| #101 asks how a failure reports the first divergent operator. WS1 answers this at operator |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Fix the Markdown heading syntax.
Add a space after the hash: # First-Divergence Localization. This resolves the MD018 warning and restores heading rendering.
🧰 Tools
🪛 markdownlint-cli2 (0.23.2)
[warning] 300-300: No space after hash on atx style heading
(MD018, no-missing-space-atx)
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` at line 300, Fix the Markdown
heading syntax in the WS1 batch invariance contract by adding a space after the
hash marker so the heading renders correctly and satisfies MD018.
Source: Linters/SAST tools
| The chain is a static ordering over invariance-declared operators: | ||
|
|
||
| ```text | ||
| 0 embedding | ||
| 1 rope | ||
| 2 rms_norm | ||
| 3 attention (output + LSE) | ||
| 4 det_gemm / linear | ||
| 5 silu / swiglu | ||
| 6 lm_head | ||
| 7 linear_logp | ||
| 8 logp / batch_invariant_logp | ||
| 9 loss reductions (grpo_loss, ratio_kl) | ||
| ``` |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Define the output scope for loss reductions.
P0 requires T1, T3, and T4 invariance checks for grpo_loss and ratio_kl, while the chain labels them as loss reductions. If these operators return batch-normalized or global reductions, changing batch size or neighbor content changes the result by design. Specify per-token outputs or fixed normalization before requiring bitwise invariance.
Also applies to: 340-352
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 304 - 317, Clarify
the output contract for the loss-reduction operators grpo_loss and ratio_kl in
the invariance chain: define them as per-token outputs or specify a fixed,
batch-independent normalization before requiring T1, T3, and T4 bitwise
invariance. Apply the same clarification to the corresponding section around the
additional loss-reduction requirements.
| The suite runs every operator against the same transform set under the same fixture family. When | ||
| multiple operators fail, the report names the **lowest `chain_index`** as the localization result: | ||
|
|
||
| ```text | ||
| first divergent operator: rms_norm (chain_index 2, T4 padding_layout, bf16, cuda) | ||
| downstream also failing: attention, lm_head, logp | ||
| verdict: investigate rms_norm; downstream failures are not independently actionable | ||
| ``` | ||
|
|
||
| This is a triage ordering, not a proof of causation. An operator early in the chain that violates | ||
| invariance will propagate to everything after it, so downstream failures carry no independent | ||
| information until the earliest one is fixed. Operators are compared independently on identical | ||
| fixtures, so this ordering is a reporting convention rather than a data dependency — an operator | ||
| can fail in isolation without any upstream failure. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Keep independent operator failures visible.
The RFC says that every failure emits a record and that operators run independently, but it also makes only the lowest chain_index actionable. A higher-index operator can fail independently even when an earlier operator also fails. Keep the lowest index as the primary triage verdict, but retain and classify all failure records.
Also applies to: 485-500
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 319 - 332, Update
the reporting contract for the failure records and primary triage verdict so
every independently failing operator remains visible and classified, including
higher-chain-index failures when an earlier operator also fails. Keep the lowest
chain_index as the primary investigation result, but explicitly distinguish
downstream failures caused by propagation from independently observed operator
failures; apply the same rule to the related section around the failure-record
requirements.
| Gaps this RFC closes: | ||
|
|
||
| 1. **No shared axis vocabulary.** The transform classes above exist today as hand-written tests | ||
| repeated across roughly thirty functions in a dozen files, with the phrase "Axis A — batch | ||
| invariance, bitwise" copy-pasted as a comment in at least five of them. There is no way to answer | ||
| "which (operator x transform x backend) cells are green?" without grepping. | ||
| 2. **The declared threshold is unwired.** `batch_invariance: {atol: 0.0, rtol: 0.0}` is in | ||
| `tolerance_contract.json` and read by no code path. | ||
| 3. **`grpo_loss`, `ratio_kl`, and sampling have no invariance tests at all**, despite being named | ||
| in #101 and sitting at the end of the chain where drift is amplified into the policy ratio. | ||
| 4. **No drift report and no localization.** Failures surface as bare pytest assertions. | ||
|
|
||
| Note that the substance is largely present and the coverage is real. What is missing is a declared | ||
| matrix, a wired threshold, and a report format — this is consolidation work more than new | ||
| verification. |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Label pending implementation as proposed.
The follow-up table shows that the generator, runner, report, and CLI are pending, but “Gaps this RFC closes” and the command sections read as current functionality. Since this PR is an RFC without implementation, use “identifies” and label the commands as proposed or post-implementation.
Also applies to: 452-477
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 397 - 411, Update
the RFC language to distinguish planned work from existing functionality: change
“Gaps this RFC closes” to describe gaps the RFC identifies, and mark the
generator, runner, report, and CLI command sections as proposed or
post-implementation. Apply the same terminology consistently in the referenced
follow-up table and command sections.
| Two additions to existing structures: | ||
|
|
||
| ```text | ||
| OperatorSpec.invariance: InvarianceDeclaration | None | ||
| chain_index: int | ||
| transforms: frozenset[str] # which classes apply | ||
| backward: bool | ||
| exempt_reason: str | None # required when invariance is None | ||
| ``` | ||
|
|
||
| An operator with `invariance=None` must state why. `NativeGemmOp` declares | ||
| `exempt_reason="non-deterministic reference baseline"`. This makes exemption a deliberate, reviewed | ||
| act rather than an omission, and it makes the matrix self-describing: every registered operator is | ||
| either covered or explicitly excused. |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Represent exemptions as a valid OperatorSpec state.
invariance can be None, but exempt_reason is shown under the declaration and is required when invariance is None. That state cannot carry the reason. Put exempt_reason on OperatorSpec, or define a discriminated union for declared and exempt operators.
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 431 - 444, Move
exempt_reason into the OperatorSpec state, or model OperatorSpec as a
discriminated union containing either an invariance declaration or an exemption
reason. Ensure operators with invariance=None can carry and require the
exemption reason, while covered operators retain their existing invariance
fields.
| [#83]: https://github.com/RL-Align/RL-Kernel/issues/83 | ||
| [#96]: https://github.com/RL-Align/RL-Kernel/issues/96 | ||
| [#101]: https://github.com/RL-Align/RL-Kernel/issues/101 |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Remove unused link reference definitions.
Reference definitions #83 and #101 are unused because those links use inline URLs. Remove the definitions or convert the corresponding links to reference-style links.
🧰 Tools
🪛 markdownlint-cli2 (0.23.2)
[warning] 537-537: Link and image reference definitions should be needed
Unused link or image reference definition: "#83"
(MD053, link-image-reference-definitions)
[warning] 539-539: Link and image reference definitions should be needed
Unused link or image reference definition: "#101"
(MD053, link-image-reference-definitions)
🤖 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 `@docs/design/ws1_batch_invariance_contract.md` around lines 537 - 539, Remove
the unused Markdown reference definitions for `#83` and `#101` from the document,
while preserving the existing inline URLs and leaving the still-referenced `#96`
definition unchanged.
Source: Linters/SAST tools
|
cc @maxiaosong1124 @frank-2077 PTAL, Do we need this PR? |
Summary
Drafts the RFC requested by #101: the batch-invariance contract for RL-Kernel's train-inference consistency suite.
Filed as WS1 rather than a new workstream — WS1 in #83 already is "full batch-invariant forward chain." This is the missing contract document over WS1's existing implementation issues, plus the localisation half of WS4.
The core proposal
Three comparison relations, so the threshold for any check follows from its kind rather than from case-by-case negotiation:
The argument for bitwise on I: the kernel, dtype, data and device are all held fixed and only the surrounding shape moves, so there is no legitimate source of floating-point difference. A tolerance there would hide exactly the reduction-schedule defect the suite exists to find. WS2 correctly refuses this stance for cross-engine comparison; WS1 can demand it.
Consequence worth noting: an I-relation needs no reference implementation — the kernel is its own oracle. That is what makes the P0 suite single-device, seconds-long, and viable as a PR gate where WS2 structurally is not.
Two findings that shaped the document
tolerance_contract.jsonalready declares"batch_invariance": {"atol": 0.0, "rtol": 0.0}, and nothing reads it —op_checks.py::_resolve_toleranceonly consultscontract["accuracy"]. Five test files independently carry a copy-pastedAxis A -- batch invariance, bitwisecomment. The policy is settled in practice; this RFC ratifies and wires it rather than proposing it.run_operator_suiteimplements the A-relation shape (two functions, one input). The I-relation is one function, two input layouts, exact comparison — different enough to warrant a siblingrun_invariance_suiterather than a flag.Also included
source_classvalues compatible with the WS2KnobDefinitionmodel so a WS1 axis lifts into a WS2 ablation without redefinition.InvarianceDeclaration | NoneonOperatorSpecwith a requiredexempt_reasonwhen null (NativeGemmOp-> "non-deterministic reference baseline"), making the matrix self-describing.grpo_lossandratio_klhave none.Positioning
WS1 is a precondition for WS2, not a competitor. The WS2 ablation ladder starts at
A0: fully aligned reference; if the underlying kernels are not batch-invariant, A0 is nondeterministic under batching and everything above it is noise.Open question for reviewers
The RFC asserts bitwise as the P0 requirement rather than leaving it open. That is the decision most worth challenging: if invariance is tolerance-based instead, the suite collapses into a WS2 knob and the separate PR-blocking CI lane loses its justification. The evidence says bitwise is already what the codebase does.
No implementation in this PR — the roadmap table lists 12 follow-up work items, deliberately left unfiled so maintainers can decide how to slice them.
Closes #101 as an RFC deliverable.
🤖 Generated with Claude Code
Summary by CodeRabbit