[WS2] Add TP-aware logprob contract and dispatch metadata - #259
[WS2] Add TP-aware logprob contract and dispatch metadata#259ryankert01 wants to merge 6 commits into
Conversation
📝 WalkthroughWalkthroughChangesWS2 logprob dispatch
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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 |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (8)
rl_engine/kernels/registry.py (3)
558-563: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDelegate
_platformto_platform_for_device.Lines 559-563 duplicate the
device is Nonebranch of_platform_for_deviceat Lines 586-590 exactly. If the ROCm or CUDA detection rule changes, one site can be updated and the other missed, and WS2 dispatch would then resolve a different platform than legacyget_op.♻️ Proposed refactor
def _platform(self) -> str: - if device_ctx.is_rocm: - return "rocm" - if device_ctx.device_type == "cuda": - return "cuda" - return "cpu" + return self._platform_for_device(None)🤖 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/registry.py` around lines 558 - 563, Update the _platform method to delegate platform resolution to _platform_for_device using the current device context, removing its duplicated ROCm/CUDA/CPU detection logic. Preserve the existing platform selection behavior so WS2 dispatch remains consistent with legacy get_op resolution.
565-582: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse
_get_or_create_backendinget_op.This helper reproduces the per-backend body of the
get_oploop at Lines 435-452: cache lookup, failed-backend skip, load, instantiate, record failure. The two copies must stay in step, or WS2 dispatch and legacy dispatch will cache and blacklist backends differently.The behavior is identical, so
get_opcan call the helper.♻️ Proposed refactor for Lines 435-452
for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] - - if backend.name in self._failed_backends: - continue - - op_class = self._load_backend(backend) - if op_class: - try: - op_instance = op_class() - self._instance_cache[backend.name] = op_instance - return op_instance - except Exception as e: - logger.error(f"Failed to instantiate {backend.name}: {e}") - self._failed_backends.add(backend.name) - else: - self._failed_backends.add(backend.name) + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance🤖 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/registry.py` around lines 565 - 582, Update get_op to call _get_or_create_backend for each backend instead of duplicating cache lookup, failed-backend filtering, loading, instantiation, and failure recording. Preserve get_op’s existing dispatch behavior by using the helper’s returned instance and continuing when it returns None.
526-534: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a dedicated dispatch error type.
Dispatch failure raises a bare
RuntimeError. Callers cannot distinguish "no backend satisfies the contract" from any other runtime failure inside the registry, so they must match on the message text. Input validation on Lines 476 and 478 already uses the typedLogprobContractError.A
LogprobDispatchError(RuntimeError)subclass keeps the existingpytest.raises(RuntimeError)assertions intests/test_logprob_contract.pypassing and gives WS2 callers a precise except clause.🤖 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/registry.py` around lines 526 - 534, The dispatch failure in the registry should use a dedicated typed exception instead of bare RuntimeError. Define or reuse LogprobDispatchError as a RuntimeError subclass, then raise it in the no-supported-backend path that builds the rejection details, while preserving the existing error message and RuntimeError compatibility.tests/test_logprob_contract.py (3)
396-402: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert decoupling in both directions.
The registry comment at
rl_engine/kernels/registry.pyLines 339-341 claims isolation in both directions. This test only proves that a WS2 candidate insertion does not reach the legacy list. Add the reverse assertion so a future change that shares one list object fails here.💚 Proposed addition
legacy = registry._priority_map[platform]["batch_invariant_logp"] assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_NATIVE_MATMUL) + assert OpBackend.PYTORCH_NATIVE_MATMUL not in registry._logprob_candidates[platform]🤖 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_logprob_contract.py` around lines 396 - 402, Extend test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map to mutate the legacy priority list after the existing assertion and verify OpBackend.PYTORCH_NATIVE does not appear in registry._logprob_candidates[platform]. Preserve the current forward-direction assertion so the test validates isolation in both directions.
29-31: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winHandle a non-divisible vocabulary in
_even_bounds.
shard = padded_vocab // tp_world_sizedrops the remainder, so the last bound ends attp_world_size * shard. The current sweep uses TP 1, 2 and 4, and 152064 divides evenly, so the helper is correct today. If a later PR extends the sweep to a TP degree that does not divide the padded vocabulary, the helper silently produces bounds thatShardingSpecrejects with "cover padded_vocab_size exactly", and the failure looks like a contract bug rather than a fixture bug.Pin the last bound to
padded_vocab.♻️ Proposed refactor
def _even_bounds(padded_vocab: int, tp_world_size: int) -> tuple[tuple[int, int], ...]: shard = padded_vocab // tp_world_size - return tuple((rank * shard, (rank + 1) * shard) for rank in range(tp_world_size)) + return tuple( + (rank * shard, padded_vocab if rank == tp_world_size - 1 else (rank + 1) * shard) + for rank in range(tp_world_size) + )🤖 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_logprob_contract.py` around lines 29 - 31, Update _even_bounds so the final shard’s upper bound is explicitly padded_vocab rather than relying on tp_world_size * shard. Preserve the existing evenly sized lower bounds and returned tuple structure while ensuring all bounds cover the full vocabulary for non-divisible TP sizes.
274-278: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider a public registration seam for logprob capabilities.
Seven tests configure dispatch by assigning to
registry._logprob_candidatesandregistry._logprob_capabilitiesand by callingregistry._platform(). The design doc states that the later deterministic vocab-parallel reference "becomes selectable by registering a capability", so production code will need the same seam.A public
register_logprob_backend(platform, backend, capability)method would let the tests stop reaching into private state and would give PR3 a supported registration path.🤖 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_logprob_contract.py` around lines 274 - 278, Add a public KernelRegistry.register_logprob_backend(platform, backend, capability) method that records the backend candidate for the platform and stores its capability, reusing the existing registration state used by _logprob_candidates and _logprob_capabilities. Update the affected tests to call this method instead of mutating private fields, while retaining _platform() only where no public platform accessor exists.rl_engine/kernels/logprob_contract.py (2)
399-402: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the implementation-kind vocabulary from one constant.
The same three kinds appear as literals here, inside
RESERVED_DISPATCH_POLICIESat Line 32, and again inKernelRegistry._logprob_policy_mismatch(registry.py Line 544). If a fourth kind is added later, all three sites must change together, and a missed site makes the kind either unselectable or unvalidated.Consider exporting a single
IMPLEMENTATION_KINDSfrozenset and buildingRESERVED_DISPATCH_POLICIESand the registry policy check from it.♻️ Proposed refactor
-RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +IMPLEMENTATION_KINDS = frozenset({"production", "reference", "deterministic"}) +RESERVED_DISPATCH_POLICIES = frozenset({"auto"}) | IMPLEMENTATION_KINDS- if self.implementation_kind not in {"production", "reference", "deterministic"}: + if self.implementation_kind not in IMPLEMENTATION_KINDS: raise LogprobContractError( - "implementation_kind must be production, reference, or deterministic" + "implementation_kind must be one of: " + + ", ".join(sorted(IMPLEMENTATION_KINDS)) )Then import
IMPLEMENTATION_KINDSinrl_engine/kernels/registry.pyand use it at Line 544.🤖 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/logprob_contract.py` around lines 399 - 402, Define and export a single IMPLEMENTATION_KINDS frozenset containing the supported implementation kinds, then use it to build RESERVED_DISPATCH_POLICIES and replace the literal membership check in LogprobContract validation. Import and reuse IMPLEMENTATION_KINDS in KernelRegistry._logprob_policy_mismatch so all vocabulary checks derive from the same constant.
345-350: 🚀 Performance & Scalability | 🔵 Trivial | ⚖️ Poor tradeoffConsider omitting the full
active_maskfrom provenance.
to_dict()copies every per-token boolean into the provenance dict.KernelRegistry.get_logprob_opembeds this dict in the provenance of each successful dispatch and callsto_dict()again to build the failure message. For realistic sequence lengths the mask dominates the payload size of a structure that exists for logging and serialization.
num_tokensandactive_token_countare already present. Consider dropping the raw mask, or emitting a compact form such as active index ranges.♻️ Proposed change
mask = { "num_tokens": self.mask.num_tokens, "active_token_count": self.mask.active_token_count, - "active_mask": list(self.mask.active_mask), "ignore_index": self.mask.ignore_index, }🤖 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/logprob_contract.py` around lines 345 - 350, Update the mask serialization in the relevant to_dict() implementation to omit the full active_mask array and retain the existing num_tokens and active_token_count fields; if mask detail is required, replace it with a compact representation such as active index ranges, ensuring KernelRegistry.get_logprob_op provenance and failure messages no longer duplicate per-token booleans.
🤖 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/ws2-tp-logprob-contract.md`:
- Around line 128-131: Revise the TP invariance claim in the logprob contract:
fixed shard-index merge order guarantees deterministic results only for a given
TP degree, not bitwise equality between TP=2 and TP=1. Either specify the
additional fixed per-shard tile decomposition required for cross-TP bitwise
equality, or describe cross-degree results as tolerance-based and defer to the
referenced `#108` tolerance table.
In `@rl_engine/kernels/logprob_contract.py`:
- Around line 484-501: Sort the `__all__` entries in `logprob_contract.py`
according to RUF022/isort ordering: place `RESERVED_DISPATCH_POLICIES` first and
ensure `LogprobDType` precedes `LogprobDispatchResult`, while preserving all
existing exports.
- Around line 379-384: Normalize backend_id by storing its stripped value after
validating it in the contract initialization flow, or reject values with leading
or trailing whitespace. Update the logic around the backend_id validation and
reserved-policy check so registry lookups and _logprob_policy_mismatch
comparisons use the same normalized identifier.
---
Nitpick comments:
In `@rl_engine/kernels/logprob_contract.py`:
- Around line 399-402: Define and export a single IMPLEMENTATION_KINDS frozenset
containing the supported implementation kinds, then use it to build
RESERVED_DISPATCH_POLICIES and replace the literal membership check in
LogprobContract validation. Import and reuse IMPLEMENTATION_KINDS in
KernelRegistry._logprob_policy_mismatch so all vocabulary checks derive from the
same constant.
- Around line 345-350: Update the mask serialization in the relevant to_dict()
implementation to omit the full active_mask array and retain the existing
num_tokens and active_token_count fields; if mask detail is required, replace it
with a compact representation such as active index ranges, ensuring
KernelRegistry.get_logprob_op provenance and failure messages no longer
duplicate per-token booleans.
In `@rl_engine/kernels/registry.py`:
- Around line 558-563: Update the _platform method to delegate platform
resolution to _platform_for_device using the current device context, removing
its duplicated ROCm/CUDA/CPU detection logic. Preserve the existing platform
selection behavior so WS2 dispatch remains consistent with legacy get_op
resolution.
- Around line 565-582: Update get_op to call _get_or_create_backend for each
backend instead of duplicating cache lookup, failed-backend filtering, loading,
instantiation, and failure recording. Preserve get_op’s existing dispatch
behavior by using the helper’s returned instance and continuing when it returns
None.
- Around line 526-534: The dispatch failure in the registry should use a
dedicated typed exception instead of bare RuntimeError. Define or reuse
LogprobDispatchError as a RuntimeError subclass, then raise it in the
no-supported-backend path that builds the rejection details, while preserving
the existing error message and RuntimeError compatibility.
In `@tests/test_logprob_contract.py`:
- Around line 396-402: Extend
test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map to mutate the
legacy priority list after the existing assertion and verify
OpBackend.PYTORCH_NATIVE does not appear in
registry._logprob_candidates[platform]. Preserve the current forward-direction
assertion so the test validates isolation in both directions.
- Around line 29-31: Update _even_bounds so the final shard’s upper bound is
explicitly padded_vocab rather than relying on tp_world_size * shard. Preserve
the existing evenly sized lower bounds and returned tuple structure while
ensuring all bounds cover the full vocabulary for non-divisible TP sizes.
- Around line 274-278: Add a public
KernelRegistry.register_logprob_backend(platform, backend, capability) method
that records the backend candidate for the platform and stores its capability,
reusing the existing registration state used by _logprob_candidates and
_logprob_capabilities. Update the affected tests to call this method instead of
mutating private fields, while retaining _platform() only where no public
platform accessor exists.
🪄 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 Plus
Run ID: 463f7688-b76b-40bb-a7f9-70cfaa95c13e
📒 Files selected for processing (7)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/design/ws2-tp-logprob-contract.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/registry.pytests/test_logprob_contract.py
Implements PR 1 of issue RL-Align#241: a typed contract for vocab-parallel selected-token logprob, mirroring the WS2 attention contract pattern. - rl_engine/kernels/logprob_contract.py: LogprobContract, ShardingSpec (per-rank vocab shard bounds, padded-vs-real vocab, TP/CP rank metadata, owner_rank resolution), MaskSpec (active-token mask, ignore_index), ReductionSpec (fp32 (max, sumexp) merge in fixed global vocab-shard index order, all-gather transport, CP declared a non-merge axis), and LogprobBackendCapability. - KernelRegistry.get_logprob_op(contract): contract-aware dispatch that only selects backends with a declared capability; incompatible or undeclared candidates are rejected with explicit reasons and never used as a silent fallback. Existing WS1 batch-invariant logp backends are declared truthfully as single-shard references, so strict WS2 requests fail loudly until the deterministic vocab-parallel TP reference (PR 3) lands. Legacy get_op() behavior is unchanged. - Design doc, runtime-dispatch and operator doc updates, and CPU-safe contract/dispatch tests covering the Qwen3-8B TP=2 BF16 target and the TP=1/2/4 sweep shapes. Tolerance values remain owned by RL-Align#108.
- docs: correct the TP-invariance claim — fixed merge order gives determinism per TP degree; cross-degree bitwise equality additionally requires a TP-degree-independent local tile decomposition (PR 3 obligation), otherwise RL-Align#108 tolerances apply - contract: store backend_id stripped so id-based dispatch matches; summarize the active mask in to_dict() provenance instead of copying every per-token boolean; sort __all__ per RUF022 - registry: add public register_logprob_backend() seam for PR 3 and tests; delegate _platform() to _platform_for_device(None); reuse _get_or_create_backend() in get_op so WS2 and legacy dispatch share one cache/blacklist code path - tests: use the registration seam instead of poking private state, pin _even_bounds' last bound for non-divisible vocabularies, assert candidate-list decoupling in both directions, cover registration replace semantics and backend_id normalization
a003a04 to
cdc11ba
Compare
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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/ws2-tp-logprob-contract.md`:
- Around line 130-136: Revise the cross-degree equality requirement in the
deterministic reference implementation discussion: do not state that
TP-degree-independent local tile decomposition alone guarantees bitwise
equality, since shard-level grouping can still differ. Require a globally
TP-partition-independent tile merge order, or describe cross-degree comparisons
as tolerance-based only; retain the existing per-degree determinism claim.
- Around line 117-125: Update the local reduction contract around local_logits
to mask columns whose global IDs are in [real_vocab_size, padded_vocab_size) to
-inf before computing each rank’s local_max and local_sumexp. Preserve the
existing fp32 computation, all-gather transport, and fixed global vocab-shard
merge order.
In `@rl_engine/kernels/registry.py`:
- Around line 463-465: Update the capability storage in the registration flow
around resolved_platform and _logprob_capabilities so entries are keyed by both
platform and backend, preventing registrations for one platform from overwriting
another. Update the lookup at the corresponding dispatch logic near line 506 to
use the same platform-scoped mapping, and add a test registering distinct
capabilities for one backend enum on two platforms.
🪄 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 Plus
Run ID: c097097d-d4e1-485a-903d-bf85b02dc456
📒 Files selected for processing (7)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/design/ws2-tp-logprob-contract.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/registry.pytests/test_logprob_contract.py
🚧 Files skipped from review as they are similar to previous changes (4)
- .github/workflows/ci.yml
- docs/operators/batch-invariant-logp.md
- docs/design/runtime-dispatch.md
- rl_engine/kernels/logprob_contract.py
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
rl_engine/kernels/logprob_contract.py (2)
401-404: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDerive the
implementation_kindvalues fromRESERVED_DISPATCH_POLICIES.The literal set at Line 401 is exactly
RESERVED_DISPATCH_POLICIESminus"auto". The two definitions must stay in sync. If a policy keyword is added to Line 32, this check will silently reject it as animplementation_kind, and dispatch by that policy will never match a backend.♻️ Proposed refactor
+# Policies that a backend can declare as its own implementation kind; "auto" +# is a selection strategy, not an implementation kind. +IMPLEMENTATION_KINDS = RESERVED_DISPATCH_POLICIES - {"auto"}- if self.implementation_kind not in {"production", "reference", "deterministic"}: + if self.implementation_kind not in IMPLEMENTATION_KINDS: raise LogprobContractError( - "implementation_kind must be production, reference, or deterministic" + "implementation_kind must be one of: " + f"{', '.join(sorted(IMPLEMENTATION_KINDS))}" )🤖 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/logprob_contract.py` around lines 401 - 404, Update the validation around implementation_kind to derive its allowed values from RESERVED_DISPATCH_POLICIES, excluding only the "auto" policy, instead of maintaining a separate literal set. Preserve the existing LogprobContractError message and validation behavior for unsupported values.
229-250: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider a compact mask representation for large batches.
active_maskstores one Pythonboolobject per token, and Line 241 and Line 250 each walk the whole sequence. For long training batches this makes contract construction O(num_tokens) in pure Python on the logprob call path.If contracts are built per invocation, accept an optional precomputed
active_token_countor a tensor-backed mask, and keep the tuple form for tests only. If contracts are built once per configuration, no change is needed.🤖 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/logprob_contract.py` around lines 229 - 250, The Logprob contract currently materializes and scans the entire active_mask during __post_init__. Update the contract construction path around __post_init__ to support a compact tensor-backed mask or optional precomputed active_token_count for large batches, while retaining tuple masks for tests and preserving validation of mask length and boolean semantics; avoid redundant full-sequence Python scans when the count is supplied.
🤖 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/logprob_contract.py`:
- Around line 387-390: Update the roles and dtypes normalization in the contract
validation method around _enum_value and frozenset to catch TypeError raised
when self.roles or self.dtypes is non-iterable, then re-raise it as
LogprobContractError consistent with active_mask and tp_world_sizes validation.
Preserve the existing enum conversion and empty-set validation behavior for
iterable inputs.
In `@rl_engine/kernels/registry.py`:
- Around line 442-464: Validate resolved_platform in register_logprob_backend
against the supported keys of self._priority_map before calling
_logprob_candidates.setdefault. Raise LogprobContractError for unknown platform
values, and preserve registration behavior for valid platforms.
---
Nitpick comments:
In `@rl_engine/kernels/logprob_contract.py`:
- Around line 401-404: Update the validation around implementation_kind to
derive its allowed values from RESERVED_DISPATCH_POLICIES, excluding only the
"auto" policy, instead of maintaining a separate literal set. Preserve the
existing LogprobContractError message and validation behavior for unsupported
values.
- Around line 229-250: The Logprob contract currently materializes and scans the
entire active_mask during __post_init__. Update the contract construction path
around __post_init__ to support a compact tensor-backed mask or optional
precomputed active_token_count for large batches, while retaining tuple masks
for tests and preserving validation of mask length and boolean semantics; avoid
redundant full-sequence Python scans when the count is supplied.
🪄 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 Plus
Run ID: a8910713-590c-4bb7-9f3b-0cd3239cbf28
📒 Files selected for processing (7)
.github/workflows/ci.ymldocs/design/runtime-dispatch.mddocs/design/ws2-tp-logprob-contract.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/registry.pytests/test_logprob_contract.py
🚧 Files skipped from review as they are similar to previous changes (4)
- docs/operators/batch-invariant-logp.md
- .github/workflows/ci.yml
- docs/design/ws2-tp-logprob-contract.md
- docs/design/runtime-dispatch.md
- docs: state that cross-TP bitwise equality needs a global tile-level merge structure independent of TP partitioning (per-shard tiles alone leave different grouping at shard boundaries), and that padded columns are masked to -inf before the local (max, sumexp) partials - registry: scope logprob capabilities per platform so the same backend enum can declare different support on cuda/rocm/cpu; validate the platform argument of register_logprob_backend against known platforms - contract: derive IMPLEMENTATION_KINDS from RESERVED_DISPATCH_POLICIES and use it for the kind check; wrap non-iterable roles/dtypes in LogprobContractError for consistent error handling - tests: cover per-platform capability scoping, unknown-platform rejection, and non-iterable roles/dtypes
…typed contract Address external review: the cross-TP bitwise guarantee lived only in prose, so a fixed-topology-deterministic backend could pass dispatch as fully conformant. - DeterminismScope (fixed_topology | cross_tp_bitwise): requested via ReductionSpec (default cross_tp_bitwise, the RL-Align#241 PR 3 target), declared per backend via determinism_scopes, enforced by dispatch; replaces the deterministic_tp_merge bool - MaskMode (explicit_active_mask | ignore_index) replaces supports_inactive_tokens: the contract permits inactive targets that do not hold ignore_index, so ignore-index-only backends are rejected for contracts with inactive tokens - LogprobOutputSpec pins the output surface: fp32 selected logprob and fp32 vocab LSE, replicated across the TP group - implementation_kind is now a tier (reference | production); determinism is no longer conflated with it, and requesting "deterministic" as a policy raises a loud error pointing at determinism_scope - fallback provenance: policy evaluation now precedes capability checks, so a candidate excluded by the caller's own policy never counts as a fallback even when it also lacks capabilities - docs: define the (-inf, 0) identity partial for padding-only or all--inf shards; document that requested_backend="auto" is not distributed-safe and specify the preflight fingerprint agreement - LogprobContract.cross_rank_fingerprint(): rank-independent identity for that preflight; provenance now records active_mask_sha256 so masks with equal active counts remain distinguishable
Fold the normative reduction semantics (padded-column masking, fp32 (max, sumexp) merge formulas, the (-inf, 0) identity partial, and the cross-TP tile-structure requirement) into the ReductionSpec and DeterminismScope docstrings, and repoint the runtime-dispatch and batch-invariant-logp doc references at the module. The contract summary moves to the PR description.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
rl_engine/kernels/logprob_contract.py (1)
527-573: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd a non-empty check for
determinism_scopes.
rolesanddtypesraiseLogprobContractErrorwhen empty (Line 545-546).determinism_scopeshas no equivalent check, and its default isfrozenset()(Line 529).
incompatibilities()unconditionally evaluatescontract.reduction.determinism_scope not in self.determinism_scopesfor every contract, so a capability left at the defaultdeterminism_scopesvalue constructs successfully but can never satisfy any contract. A developer who forgets to declaredeterminism_scopeswhen registering a new backend gets no error at registration time. They only see repeated "determinism_scope=... is unsupported" rejections at dispatch time, with no direct signal that the field was never set.Add a construction-time check so this failure mode surfaces immediately, consistent with the roles/dtypes check.
🐛 Proposed fix
except TypeError as exc: raise LogprobContractError( "mask_modes and determinism_scopes must be iterables of enum values" ) from exc + if not determinism_scopes: + raise LogprobContractError( + "determinism_scopes must not be empty; declare at least one " + "DeterminismScope this backend actually guarantees" + ) for flag_name in ("supports_vocab_padding", "exports_vocab_lse"):🤖 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/logprob_contract.py` around lines 527 - 573, Add determinism_scopes to the non-empty validation in BackendCapability.__post_init__, alongside roles and dtypes, so the default empty frozenset raises LogprobContractError during construction. Preserve the existing enum normalization and error behavior for non-empty values.
🤖 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/logprob_contract.py`:
- Around line 496-515: Update get_logprob_op in registry.py to reject
requested_backend="auto" when contract.sharding.tp_world_size > 1, preventing
dispatch to TP=1-only candidates. Allow this path only when the caller
explicitly indicates that cross-rank fingerprint and resolved-backend agreement
has already been preflighted, reusing the existing cross_rank_fingerprint
contract and established preflight state if available.
---
Outside diff comments:
In `@rl_engine/kernels/logprob_contract.py`:
- Around line 527-573: Add determinism_scopes to the non-empty validation in
BackendCapability.__post_init__, alongside roles and dtypes, so the default
empty frozenset raises LogprobContractError during construction. Preserve the
existing enum normalization and error behavior for non-empty values.
🪄 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 Plus
Run ID: 3f6fb022-c2a7-4f6d-9eea-7efe717732fc
📒 Files selected for processing (5)
docs/design/runtime-dispatch.mddocs/operators/batch-invariant-logp.mdrl_engine/kernels/logprob_contract.pyrl_engine/kernels/registry.pytests/test_logprob_contract.py
🚧 Files skipped from review as they are similar to previous changes (3)
- docs/operators/batch-invariant-logp.md
- docs/design/runtime-dispatch.md
- rl_engine/kernels/registry.py
| def cross_rank_fingerprint(self) -> str: | ||
| """Rank-independent identity for preflight agreement across ranks. | ||
|
|
||
| Excludes ``tp_rank``/``cp_rank`` (and their derived local bounds) so | ||
| every rank of one logical invocation computes the same value. | ||
| All-gathering this fingerprint together with the resolved backend id | ||
| and aborting on mismatch is the documented preflight for distributed | ||
| dispatch; ``requested_backend="auto"`` is not distributed-safe | ||
| without it. | ||
| """ | ||
|
|
||
| payload = self.to_dict() | ||
| payload["sharding"] = { | ||
| key: value | ||
| for key, value in payload["sharding"].items() | ||
| if key not in {"tp_rank", "cp_rank", "local_vocab_start", "local_vocab_end"} | ||
| } | ||
| encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")) | ||
| return hashlib.sha256(encoded.encode("utf-8")).hexdigest() | ||
|
|
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check whether get_logprob_op enforces auto+TP>1 rejection.
rg -n -B3 -A15 'def get_logprob_op' rl_engine/kernels/registry.py
rg -n 'auto' rl_engine/kernels/registry.py
rg -n 'requested_backend.*auto|auto.*tp_world_size' -i rl_engine/kernels/registry.pyRepository: RL-Align/RL-Kernel
Length of output: 1479
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== registry.py relevant sections =="
sed -n '110,175p' rl_engine/kernels/registry.py
echo
sed -n '490,620p' rl_engine/kernels/registry.py
echo
echo "== LogprobContract and cross_rank_fingerprint implementations/usages =="
sed -n '440,525p' rl_engine/kernels/logprob_contract.py
rg -n 'cross_rank_fingerprint|requested_backend|get_logprob_op|tp_world_size|WorldSize' rl_engine/kernels -SRepository: RL-Align/RL-Kernel
Length of output: 16153
Enforce the auto + TP>1 preflight contract.
get_logprob_op(requested_backend="auto") currently accepts candidates with tp_world_sizes=(1,), so it can dispatch to TP=1-only backends even when contract.sharding.tp_world_size > 1. The cross_rank_fingerprint() docstring promises rank-independent preflight protection for this case, but no caller check is present in registry.py. Add an explicit rejection for requested_backend="auto" when tp_world_size > 1, unless the caller has already preflighted fingerprint/backend agreement.
🧰 Tools
🪛 ast-grep (0.45.0)
[info] 512-512: use jsonify instead of json.dumps for JSON output
Context: json.dumps(payload, sort_keys=True, separators=(",", ":"))
Note: [CWE-116] Improper Encoding or Escaping of Output.
(use-jsonify)
🤖 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/logprob_contract.py` around lines 496 - 515, Update
get_logprob_op in registry.py to reject requested_backend="auto" when
contract.sharding.tp_world_size > 1, preventing dispatch to TP=1-only
candidates. Allow this path only when the caller explicitly indicates that
cross-rank fingerprint and resolved-backend agreement has already been
preflighted, reusing the existing cross_rank_fingerprint contract and
established preflight state if available.
Shrink class docstrings toward the attention-contract one-liner style and cut design-rationale comments; the normative reduction semantics stay in the ReductionSpec and DeterminismScope docstrings.
1 participant
Part of #241 — PR 1: typed contract + dispatch metadata for TP-aware logprob. Spec only; kernels land in PR 3.
logprob_contract.py—ShardingSpec,MaskSpec,ReductionSpec,LogprobOutputSpec,LogprobBackendCapability; invalid metadata fails loudly at construction.KernelRegistry.get_logprob_op(contract)— capability-gated dispatch with explicit rejection reasons, no silent fallback. WS1 backends are declared as TP=1, so strict TP requests fail until PR 3's backend registers. Legacyget_op()untouched.Reduction semantics (fp32
(max, sumexp)partials, all-gather transport, fixed merge order,DeterminismScope) are in theReductionSpec/DeterminismScopedocstrings. Shard layout and padding arithmetic follow Megatron; padding is excluded from the LSE to match vLLM's rollout side.Out of scope: tolerances (#108), drift-report format (#116), backend invocation and preflight (PR 3/4).
Validation: 44 tests pass, black/isort/flake8/ruff/mypy clean,
mkdocs build --strictpasses.