From b7ffb3ad496535909215e43b0bb8dfea446aa0a5 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 20:03:13 +0800 Subject: [PATCH 1/6] feat(ws2): add TP-aware logprob contract and dispatch metadata Implements PR 1 of issue #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 #108. --- .github/workflows/ci.yml | 3 + docs/design/runtime-dispatch.md | 7 + docs/design/ws2-tp-logprob-contract.md | 196 ++++++++++ docs/operators/batch-invariant-logp.md | 13 + rl_engine/kernels/logprob_contract.py | 501 +++++++++++++++++++++++++ rl_engine/kernels/registry.py | 192 ++++++++++ tests/test_logprob_contract.py | 402 ++++++++++++++++++++ 7 files changed, 1314 insertions(+) create mode 100644 docs/design/ws2-tp-logprob-contract.md create mode 100644 rl_engine/kernels/logprob_contract.py create mode 100644 tests/test_logprob_contract.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 92cd0433..28cdb58d 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -76,6 +76,9 @@ jobs: run: | python -m pytest tests/test_kv_cache_attention.py -v -k "not large and not gpu" + - name: Run WS2 Logprob Contract Tests (CPU-safe) + run: python -m pytest tests/test_logprob_contract.py -v + docs: runs-on: ubuntu-latest steps: diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index bedf3475..84f29ec3 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -11,6 +11,13 @@ logical type, and the registry selects the first available backend for the curre 4. Cache successfully constructed operator instances. 5. Skip backends that already failed in the current process. +WS2 TP-aware logprob uses the stricter `KernelRegistry.get_logprob_op(contract)` path. In +addition to platform priority, this path requires a backend capability descriptor and checks +the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token +support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible +candidates produce explicit rejection reasons and are never used as an undeclared fallback. +See [WS2 TP-aware logprob contract](ws2-tp-logprob-contract.md). + ## LogP Priority | Platform | Priority | diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md new file mode 100644 index 00000000..a392393b --- /dev/null +++ b/docs/design/ws2-tp-logprob-contract.md @@ -0,0 +1,196 @@ +# WS2 TP-Aware Logprob Contract + +Status: PR1 contract and dispatch metadata + +Tracking and shared contracts: + +- [#241: TP-aware deterministic logprob](https://github.com/RL-Align/RL-Kernel/issues/241) +- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) +- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) +- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) +- [#116: WS2 tolerance and drift-report format](https://github.com/RL-Align/RL-Kernel/issues/116) +- [Cross-config logprob drift contract](ws2_cross_config_logprob_drift_contract.md) + +## Scope + +This contract describes the logical inputs and deterministic reduction semantics for +selected-token log-probability under vocab-parallel tensor parallelism (TP): + +```text +selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :]) +``` + +Under vocab-parallel TP each rank holds one vocabulary shard, so the vocabulary-wide +`logsumexp` requires a cross-rank reduction. This contract lets runtime dispatch reject a +backend whose numerical semantics do not match the requested layout. + +This PR1 layer does not shard tensors, launch a collective, merge `(max, sumexp)` partial +states, or implement a kernel. The single-GPU harness registration, the deterministic +vocab-parallel TP reference, and the cross-config integration belong to later PRs in #241. + +Context parallelism (CP) is a declared non-merge axis. CP partitions tokens, never the +vocabulary, so the logprob reduction spans TP vocab shards only. CP rank metadata is carried +for provenance and must never widen the merge. + +## Contract Objects + +`rl_engine.kernels.logprob_contract` defines: + +- `LogprobContract`: role, logits dtype, mask, sharding, reduction, and LSE export; +- `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank + metadata, and target-token ownership; +- `MaskSpec`: active-token mask and ignore index; +- `ReductionSpec`: fixed `(max, sumexp)` merge semantics; +- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports. + +Construction performs validation immediately. A structurally valid contract means that the +request is complete and internally consistent; it does not mean that an installed backend can +materialize it. + +`ShardingSpec.vocab_shard_bounds` lists every TP rank's half-open `[start, end)` vocab range +indexed by TP rank. The full table is required on every rank: it defines target ownership +and the fixed merge order without any collective, and it makes an incomplete or overlapping +partition a loud construction-time error instead of a silent runtime divergence. +`ShardingSpec.owner_rank(token_id)` resolves the unique owning rank for a real-vocab token +and rejects everything else. + +`padded_vocab_size` is the shard-covered (weight) vocabulary; `real_vocab_size` is the +tokenizer vocabulary. Padding columns occupy `[real_vocab_size, padded_vocab_size)` and must +be excluded from the logsumexp by any conforming implementation. The two sizes are equal +when the vocabulary is unpadded. + +Inactive tokens (prompt, padding, masked-out response positions) are excluded from every +drift aggregate and are exempt from the exactly-one-owner target gather; their targets may +legally hold `ignore_index`. `ignore_index` must not collide with the real vocabulary. + +## Qwen3-8B TP=2 BF16 Example + +```python +from rl_engine.kernels.logprob_contract import ( + LogprobContract, + MaskSpec, + ReductionSpec, + ShardingSpec, +) + +sharding = ShardingSpec( + tp_rank=0, + tp_world_size=2, + vocab_shard_bounds=((0, 76032), (76032, 152064)), + real_vocab_size=151936, + padded_vocab_size=152064, + cp_rank=0, + cp_world_size=2, +) + +contract = LogprobContract( + role="train", + dtype="bf16", + mask=MaskSpec( + num_tokens=8, + active_mask=(False, False, True, True, True, True, True, False), + ignore_index=-100, + ), + sharding=sharding, + reduction=ReductionSpec(), +) +``` + +Each rank owns one contiguous vocab shard; the 128 padding columns at the end of rank 1's +shard are outside the real vocabulary and never contribute to the logsumexp. The two leading +prompt tokens and the trailing padding token are inactive. + +## Reduction Semantics + +The only PR1 reduction contract is: + +```text +partial state: (local_max, local_sumexp), fp32 +merge: max_sumexp +merge_axis: tp_vocab +order: global_vocab_shard_index +transport: all_gather +downcast_at: final_write +engine: in_op_reference +``` + +Every rank computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` in +fp32, the partials travel by all-gather (collectives are transport only, never a numerical +reduction), and every rank merges in fixed global vocab-shard index order: + +```text +M = max_l(m_l) +S = sum_l(s_l * exp(m_l - M)) +LSE = M + log(S) +selected_logp = target_logit - LSE +``` + +The selected target logit comes from a masked single-owner gather: exactly one rank holds +each active token's target column. Downcast happens only at the final write. Because the +merge order is fixed by shard index, TP=2 is bitwise-equal to TP=1 by construction; averaging +per-rank logsumexp values or letting a collective reduce numerically is not conformant. + +The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift +reports follow the #116 format. This contract does not introduce another tolerance table. +The selected-token metric remains the cross-config convention: + +```text +dlogp = training-side recomputed logp - rollout-side old logp +``` + +computed over active response tokens only. + +## Contract-Aware Dispatch + +Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: + +```python +result = kernel_registry.get_logprob_op(contract) +op = result.op +provenance = result.provenance +``` + +Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, +TP/CP degree, padded-vs-real vocab masking, inactive-token support, vocab-domain LSE export, +and deterministic TP merge. An undeclared or incompatible backend is skipped with an +explicit rejection reason; there is no silent fallback. + +`requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | +`reference` | `deterministic`; default `auto`) or an exact, case-sensitive stable backend +id. Strictness comes from the contract's capability checks, not from the policy string. A +backend id may never shadow a policy keyword; capability construction rejects that. The +provenance `fallback` flag reports only capability or load rejections of otherwise-eligible +candidates — skips caused purely by the caller's own policy filter are not fallbacks. + +WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy +`batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does +not change what legacy `get_op("batch_invariant_logp")` returns to WS1 callers. + +The current WS1 batch-invariant logp implementations are single-shard (TP=1) references: +they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard +metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict +WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference +becomes selectable by registering a capability that truthfully declares those features; no +controller branch or silent fallback is required. + +Successful dispatch provenance records: + +- requested and actual backend ids; +- platform and fallback status; +- prior candidate rejection reasons; +- the complete requested contract, including shard bounds, padded and real vocab sizes, + merge semantics, and the explicit `cp_is_merge_axis: false` declaration; +- the selected backend capability descriptor. + +## Validation + +Contract and dispatch behavior are covered by: + +```bash +python -m pytest tests/test_logprob_contract.py -q +``` + +The tests include Qwen3-8B TP=2 BF16 construction with padded vocab, the TP=1/2/4 sweep +shapes, incomplete/overlapping shard-bound rejection, owner-rank resolution, active-mask and +ignore-index validation, fp32-accumulation and merge-semantics enforcement, undeclared +backend rejection, no incompatible fallback, and JSON-compatible provenance. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index fbc0e9f1..b0671c40 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -54,6 +54,19 @@ CUDA priority list when the extension exposes `_C.batch_invariant_logp_sm90` (built with `KERNEL_ALIGN_FORCE_SM90=1`) on an SM90 device. On any other build or device, dispatch is unchanged (Triton -> PyTorch). +### WS2 TP-aware dispatch + +WS2 distributed callers use a separate contract-aware entry point, +`kernel_registry.get_logprob_op(contract)`. It validates explicit vocab-shard ownership, +padded-vs-real vocab metadata, active-token masking, and fixed `(max, sumexp)` merge +semantics before selecting a backend. Legacy `get_op("batch_invariant_logp")` behavior +remains unchanged. + +The backends above are single-shard (TP=1) references and do not yet export vocab-domain +LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 +requests instead of being selected as a silent fallback. See +[WS2 TP-aware logprob contract](../design/ws2-tp-logprob-contract.md). + ## Benchmarks `benchmarks/benchmark_batch_invariant_logp.py` compares Native, Triton, and the diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py new file mode 100644 index 00000000..8f2667e1 --- /dev/null +++ b/rl_engine/kernels/logprob_contract.py @@ -0,0 +1,501 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""Typed WS2 contract for TP-aware selected-token log-probability. + +The objects in this module describe a vocab-parallel logprob invocation: + +``selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :])`` + +Under vocab-parallel tensor parallelism the vocabulary-wide ``logsumexp`` +requires cross-rank reduction. This module only *describes* that invocation +(shard ownership, merge semantics, mask/ignore-index metadata); it does not +shard tensors, launch collectives, or implement the ``(max, sumexp)`` merge. +Keeping description and materialization separate lets dispatch reject an +incompatible backend before any numerically different path is launched. + +Context parallelism is a declared non-merge axis: CP partitions tokens, never +the vocabulary, so the logprob reduction spans TP vocab shards only. CP rank +metadata is carried for provenance and must never widen the merge. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from enum import Enum +from typing import Any, TypeVar + +_EnumT = TypeVar("_EnumT", bound=Enum) + +# Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable +# backend id must never shadow one of these, or it becomes unselectable by id. +RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) + + +class LogprobContractError(ValueError): + """Raised when logprob metadata does not describe a valid invocation.""" + + +class LogprobRole(str, Enum): + TRAIN = "train" + INFER = "infer" + + +class LogprobDType(str, Enum): + BF16 = "bf16" + FP16 = "fp16" + FP32 = "fp32" + + +class LogprobMerge(str, Enum): + """Merge primitive for per-shard partial states. + + Every rank contributes ``(local_max, local_sumexp)`` computed in the + accumulation dtype; the merged result is + ``M = max(m_l)``, ``S = sum(s_l * exp(m_l - M))``, ``LSE = M + log(S)``. + """ + + MAX_SUMEXP = "max_sumexp" + + +class MergeAxis(str, Enum): + """The only reduction axis of this contract; CP is a non-merge axis.""" + + TP_VOCAB = "tp_vocab" + + +class ReductionOrder(str, Enum): + GLOBAL_VOCAB_SHARD_INDEX = "global_vocab_shard_index" + + +class ReductionTransport(str, Enum): + """Collectives move partial states only; they never reduce numerically.""" + + ALL_GATHER = "all_gather" + + +class DowncastPoint(str, Enum): + FINAL_WRITE = "final_write" + + +class ReductionEngine(str, Enum): + IN_OP_REFERENCE = "in_op_reference" + + +def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: + try: + return enum_type(value) + except (TypeError, ValueError) as exc: + allowed = ", ".join(item.value for item in enum_type) + raise LogprobContractError(f"{field} must be one of: {allowed}; got {value!r}") from exc + + +def _positive_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must be a positive integer; got {value!r}") + return value + + +def _non_negative_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int) or value < 0: + raise LogprobContractError(f"{field} must be a non-negative integer; got {value!r}") + return value + + +def _plain_int(value: Any, field: str) -> int: + if isinstance(value, bool) or not isinstance(value, int): + raise LogprobContractError(f"{field} must be an integer; got {value!r}") + return value + + +@dataclass(frozen=True) +class ShardingSpec: + """Logical vocab-parallel TP ownership for one logprob invocation. + + ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` + vocab range indexed by TP rank. The full table is required on every rank: + it defines target-token ownership and the fixed global-shard-index merge + order without any collective, and makes an incomplete partition a loud + construction-time error instead of a silent runtime divergence. + + ``padded_vocab_size`` is the shard-covered (weight) vocabulary; + ``real_vocab_size`` is the tokenizer vocabulary. Padding columns occupy + ``[real_vocab_size, padded_vocab_size)`` and must be excluded from the + logsumexp by any conforming implementation. + """ + + tp_rank: int + tp_world_size: int + vocab_shard_bounds: tuple[tuple[int, int], ...] + real_vocab_size: int + padded_vocab_size: int + cp_rank: int = 0 + cp_world_size: int = 1 + + def __post_init__(self) -> None: + tp_world_size = _positive_int(self.tp_world_size, "tp_world_size") + tp_rank = _non_negative_int(self.tp_rank, "tp_rank") + if tp_rank >= tp_world_size: + raise LogprobContractError( + f"tp_rank={tp_rank} must be smaller than tp_world_size={tp_world_size}" + ) + cp_world_size = _positive_int(self.cp_world_size, "cp_world_size") + cp_rank = _non_negative_int(self.cp_rank, "cp_rank") + if cp_rank >= cp_world_size: + raise LogprobContractError( + f"cp_rank={cp_rank} must be smaller than cp_world_size={cp_world_size}" + ) + + real_vocab_size = _positive_int(self.real_vocab_size, "real_vocab_size") + padded_vocab_size = _positive_int(self.padded_vocab_size, "padded_vocab_size") + if padded_vocab_size < real_vocab_size: + raise LogprobContractError( + f"padded_vocab_size={padded_vocab_size} must not be smaller than " + f"real_vocab_size={real_vocab_size}" + ) + + try: + bounds = tuple((pair[0], pair[1]) for pair in self.vocab_shard_bounds) + except (TypeError, IndexError) as exc: + raise LogprobContractError( + "vocab_shard_bounds must be an iterable of (start, end) integer pairs" + ) from exc + if len(bounds) != tp_world_size: + raise LogprobContractError( + "vocab_shard_bounds must declare exactly one (start, end) pair per TP rank; " + f"got {len(bounds)} pairs for tp_world_size={tp_world_size}" + ) + expected_start = 0 + for rank, (start, end) in enumerate(bounds): + start = _plain_int(start, f"vocab_shard_bounds[{rank}][0]") + end = _plain_int(end, f"vocab_shard_bounds[{rank}][1]") + if end <= start: + raise LogprobContractError( + f"vocab_shard_bounds[{rank}] must satisfy end > start; got [{start}, {end})" + ) + if start != expected_start: + raise LogprobContractError( + "vocab_shard_bounds must form a contiguous [0, padded_vocab_size) " + f"partition in TP-rank order; rank {rank} starts at {start}, " + f"expected {expected_start}" + ) + expected_start = end + if expected_start != padded_vocab_size: + raise LogprobContractError( + "vocab_shard_bounds must cover padded_vocab_size exactly; " + f"covered {expected_start}, declared {padded_vocab_size}" + ) + object.__setattr__(self, "vocab_shard_bounds", bounds) + + @property + def local_vocab_start(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][0] + + @property + def local_vocab_end(self) -> int: + return self.vocab_shard_bounds[self.tp_rank][1] + + @property + def local_vocab_size(self) -> int: + start, end = self.vocab_shard_bounds[self.tp_rank] + return end - start + + def owner_rank(self, token_id: int) -> int: + """Return the unique TP rank owning ``token_id``; error outside real vocab.""" + + token_id = _plain_int(token_id, "token_id") + if token_id < 0 or token_id >= self.real_vocab_size: + raise LogprobContractError( + f"token_id={token_id} is outside the real vocabulary " + f"[0, {self.real_vocab_size}); mask it as inactive instead" + ) + for rank, (start, end) in enumerate(self.vocab_shard_bounds): + if start <= token_id < end: + return rank + raise LogprobContractError( + f"token_id={token_id} is not covered by any declared vocab shard" + ) + + +@dataclass(frozen=True) +class MaskSpec: + """Active-token ownership for one logprob invocation. + + Inactive tokens are excluded from every drift aggregate and are exempt + from the exactly-one-owner target gather; their targets may legally hold + ``ignore_index``. + """ + + num_tokens: int + active_mask: tuple[bool, ...] + ignore_index: int = -100 + _active_token_count: int = field(init=False, repr=False, compare=False) + + def __post_init__(self) -> None: + num_tokens = _positive_int(self.num_tokens, "num_tokens") + _plain_int(self.ignore_index, "ignore_index") + try: + active_mask = tuple(self.active_mask) + except TypeError as exc: + raise LogprobContractError("active_mask must be an iterable of booleans") from exc + for index, value in enumerate(active_mask): + if not isinstance(value, bool): + raise LogprobContractError(f"active_mask[{index}] must be a bool; got {value!r}") + if len(active_mask) != num_tokens: + raise LogprobContractError( + "active_mask must contain exactly one entry per token; " + f"got {len(active_mask)} entries for num_tokens={num_tokens}" + ) + object.__setattr__(self, "active_mask", active_mask) + object.__setattr__(self, "_active_token_count", sum(active_mask)) + + @property + def active_token_count(self) -> int: + return self._active_token_count + + +@dataclass(frozen=True) +class ReductionSpec: + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics.""" + + merge: LogprobMerge = LogprobMerge.MAX_SUMEXP + merge_axis: MergeAxis = MergeAxis.TP_VOCAB + acc_dtype: LogprobDType = LogprobDType.FP32 + order: ReductionOrder = ReductionOrder.GLOBAL_VOCAB_SHARD_INDEX + transport: ReductionTransport = ReductionTransport.ALL_GATHER + downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE + engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + + def __post_init__(self) -> None: + object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) + object.__setattr__( + self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") + ) + object.__setattr__( + self, "acc_dtype", _enum_value(LogprobDType, self.acc_dtype, "acc_dtype") + ) + object.__setattr__(self, "order", _enum_value(ReductionOrder, self.order, "order")) + object.__setattr__( + self, "transport", _enum_value(ReductionTransport, self.transport, "transport") + ) + object.__setattr__( + self, "downcast_at", _enum_value(DowncastPoint, self.downcast_at, "downcast_at") + ) + object.__setattr__(self, "engine", _enum_value(ReductionEngine, self.engine, "engine")) + if self.acc_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"TP logprob accumulation must be fp32; got {self.acc_dtype.value}" + ) + + +@dataclass(frozen=True) +class LogprobContract: + """Complete semantic request consumed by contract-aware dispatch.""" + + role: LogprobRole + dtype: LogprobDType + mask: MaskSpec + sharding: ShardingSpec + reduction: ReductionSpec + export_lse: bool = True + + def __post_init__(self) -> None: + object.__setattr__(self, "role", _enum_value(LogprobRole, self.role, "role")) + object.__setattr__(self, "dtype", _enum_value(LogprobDType, self.dtype, "dtype")) + if not isinstance(self.mask, MaskSpec): + raise LogprobContractError("mask must be a MaskSpec") + if not isinstance(self.sharding, ShardingSpec): + raise LogprobContractError("sharding must be a ShardingSpec") + if not isinstance(self.reduction, ReductionSpec): + raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.export_lse, bool) or not self.export_lse: + raise LogprobContractError( + "export_lse must be True for the WS2 vocab-domain LSE drift contract" + ) + if 0 <= self.mask.ignore_index < self.sharding.real_vocab_size: + raise LogprobContractError( + f"ignore_index={self.mask.ignore_index} must not collide with the real " + f"vocabulary [0, {self.sharding.real_vocab_size})" + ) + + def to_dict(self) -> dict[str, Any]: + """Return stable, JSON-compatible requested-contract provenance.""" + + sharding = { + "tp_rank": self.sharding.tp_rank, + "tp_world_size": self.sharding.tp_world_size, + "cp_rank": self.sharding.cp_rank, + "cp_world_size": self.sharding.cp_world_size, + "vocab_shard_bounds": [list(pair) for pair in self.sharding.vocab_shard_bounds], + "real_vocab_size": self.sharding.real_vocab_size, + "padded_vocab_size": self.sharding.padded_vocab_size, + "local_vocab_start": self.sharding.local_vocab_start, + "local_vocab_end": self.sharding.local_vocab_end, + } + reduction = { + "merge": self.reduction.merge.value, + "merge_axis": self.reduction.merge_axis.value, + "acc_dtype": self.reduction.acc_dtype.value, + "order": self.reduction.order.value, + "transport": self.reduction.transport.value, + "downcast_at": self.reduction.downcast_at.value, + "engine": self.reduction.engine.value, + "cp_is_merge_axis": False, + } + 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, + } + return { + "semantic_operator": "selected_token_logprob", + "role": self.role.value, + "dtype": self.dtype.value, + "export_lse": self.export_lse, + "lse_domain": "vocab", + "mask": mask, + "sharding": sharding, + "reduction": reduction, + } + + +@dataclass(frozen=True) +class LogprobBackendCapability: + """Capabilities a concrete backend declares to contract-aware dispatch.""" + + backend_id: str + roles: frozenset[LogprobRole] + dtypes: frozenset[LogprobDType] + tp_world_sizes: tuple[int, ...] | None = None + cp_world_sizes: tuple[int, ...] | None = None + supports_vocab_padding: bool = False + supports_inactive_tokens: bool = False + exports_vocab_lse: bool = False + deterministic_tp_merge: bool = False + implementation_kind: str = "production" + + def __post_init__(self) -> None: + if not isinstance(self.backend_id, str) or not self.backend_id.strip(): + raise LogprobContractError("backend_id must be a non-empty string") + if self.backend_id.strip().lower() in RESERVED_DISPATCH_POLICIES: + raise LogprobContractError( + f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" + ) + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + if not roles or not dtypes: + raise LogprobContractError("backend roles and dtypes must not be empty") + tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") + cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") + for flag_name in ( + "supports_vocab_padding", + "supports_inactive_tokens", + "exports_vocab_lse", + "deterministic_tp_merge", + ): + if not isinstance(getattr(self, flag_name), bool): + raise LogprobContractError(f"{flag_name} must be a bool") + if self.implementation_kind not in {"production", "reference", "deterministic"}: + raise LogprobContractError( + "implementation_kind must be production, reference, or deterministic" + ) + object.__setattr__(self, "roles", roles) + object.__setattr__(self, "dtypes", dtypes) + object.__setattr__(self, "tp_world_sizes", tp_world_sizes) + object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + + @staticmethod + def _validated_world_sizes( + values: tuple[int, ...] | None, field: str + ) -> tuple[int, ...] | None: + if values is None: + return None + try: + sizes = tuple(values) + except TypeError as exc: + raise LogprobContractError(f"{field} must be an iterable of integers") from exc + if not sizes: + raise LogprobContractError(f"{field} must not be empty; use None for unrestricted") + for value in sizes: + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise LogprobContractError(f"{field} must contain positive values; got {value!r}") + if len(set(sizes)) != len(sizes): + raise LogprobContractError(f"{field} must not contain duplicates") + return sizes + + def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: + """Explain every reason this backend cannot materialize ``contract``.""" + + reasons: list[str] = [] + if contract.role not in self.roles: + reasons.append(f"role={contract.role.value} is unsupported") + if contract.dtype not in self.dtypes: + reasons.append(f"dtype={contract.dtype.value} is unsupported") + tp_size = contract.sharding.tp_world_size + cp_size = contract.sharding.cp_world_size + if self.tp_world_sizes is not None and tp_size not in self.tp_world_sizes: + reasons.append(f"TP={tp_size} is unsupported") + if self.cp_world_sizes is not None and cp_size not in self.cp_world_sizes: + reasons.append(f"CP={cp_size} is unsupported") + if ( + contract.sharding.padded_vocab_size != contract.sharding.real_vocab_size + and not self.supports_vocab_padding + ): + reasons.append("padded-vs-real vocab masking is unsupported") + if ( + contract.mask.active_token_count != contract.mask.num_tokens + and not self.supports_inactive_tokens + ): + reasons.append("inactive-token (ignore_index) masking is unsupported") + if contract.export_lse and not self.exports_vocab_lse: + reasons.append("vocab-domain LSE export is unsupported") + if tp_size > 1 and not self.deterministic_tp_merge: + reasons.append("deterministic TP (max, sumexp) merge is unsupported") + return tuple(reasons) + + def supports(self, contract: LogprobContract) -> bool: + return not self.incompatibilities(contract) + + def to_dict(self) -> dict[str, Any]: + return { + "backend_id": self.backend_id, + "roles": sorted(role.value for role in self.roles), + "dtypes": sorted(dtype.value for dtype in self.dtypes), + "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, + "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, + "supports_vocab_padding": self.supports_vocab_padding, + "supports_inactive_tokens": self.supports_inactive_tokens, + "exports_vocab_lse": self.exports_vocab_lse, + "deterministic_tp_merge": self.deterministic_tp_merge, + "implementation_kind": self.implementation_kind, + } + + +@dataclass(frozen=True) +class LogprobDispatchResult: + """A concrete backend plus the actual provenance bound to the request.""" + + op: Any + capability: LogprobBackendCapability + provenance: dict[str, Any] + + +__all__ = [ + "DowncastPoint", + "LogprobBackendCapability", + "LogprobContract", + "LogprobContractError", + "LogprobDispatchResult", + "LogprobDType", + "LogprobMerge", + "LogprobRole", + "MaskSpec", + "MergeAxis", + "RESERVED_DISPATCH_POLICIES", + "ReductionEngine", + "ReductionOrder", + "ReductionSpec", + "ReductionTransport", + "ShardingSpec", +] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index fb2feb6f..35ae951b 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -8,6 +8,14 @@ import torch +from rl_engine.kernels.logprob_contract import ( + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDispatchResult, + LogprobDType, + LogprobRole, +) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -164,6 +172,51 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() + # These descriptors report what the existing WS1 batch-invariant logp + # implementations actually support: single-shard (TP=1) logits with + # ignore-index masking, no vocab-shard metadata, no padded-vs-real + # vocab distinction, and no public vocab-domain LSE export. A strict + # WS2 request is rejected with explicit reasons until the deterministic + # vocab-parallel TP reference backend lands (issue #241 PR 3) instead + # of silently selecting an incompatible fallback. + common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) + common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) + self._logprob_capabilities = { + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="pytorch-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + implementation_kind="reference", + ), + OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( + backend_id="triton-batch-invariant-logp-ws1", + roles=common_logprob_roles, + dtypes=common_logprob_dtypes, + tp_world_sizes=(1,), + supports_vocab_padding=False, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + implementation_kind="deterministic", + ), + OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( + backend_id="cuda-batch-invariant-logp-sm90-ws1", + roles=common_logprob_roles, + dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), + tp_world_sizes=(1,), + supports_vocab_padding=False, + supports_inactive_tokens=True, + exports_vocab_lse=False, + deterministic_tp_merge=False, + implementation_kind="deterministic", + ), + } + self._priority_map = { "cuda": { "logp": [ @@ -280,6 +333,17 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() + # WS2 contract-aware dispatch owns its candidate list. It is seeded + # from the legacy batch_invariant_logp priority (after hardware/env + # adjustments) but deliberately decoupled afterwards: registering a + # TP-vocab backend for WS2 dispatch must not change what legacy + # get_op("batch_invariant_logp") returns to WS1 callers, and vice + # versa. + self._logprob_candidates: Dict[str, list] = { + platform: list(ops.get("batch_invariant_logp", [])) + for platform, ops in self._priority_map.items() + } + def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() if rocm_attn_backend in {"flash_attn", "flash-attn", "flash_attention"}: @@ -389,6 +453,134 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + def get_logprob_op( + self, + contract: LogprobContract, + *, + requested_backend: str = "auto", + ) -> LogprobDispatchResult: + """Resolve only a backend that explicitly supports the WS2 logprob contract. + + This entry point is intentionally separate from legacy ``get_op`` so + existing callers retain their current behavior while WS2 callers cannot + silently fall back to a backend with different distributed semantics. + + ``requested_backend`` is either a case-insensitive policy keyword + (``auto`` | ``production`` | ``reference`` | ``deterministic``) or an + exact, case-sensitive stable backend id. Strictness comes from the + contract's capability checks, not from this policy string, so the + default is ``auto``. + """ + + if not isinstance(contract, LogprobContract): + raise LogprobContractError("contract must be a LogprobContract") + if not isinstance(requested_backend, str) or not requested_backend.strip(): + raise LogprobContractError("requested_backend must be a non-empty string") + requested_backend = requested_backend.strip() + + platform = self._platform() + candidates = self._logprob_candidates.get(platform, []) + rejected: list[str] = [] + # provenance["fallback"] reports only capability/load rejections of + # otherwise-eligible candidates; skips caused purely by the caller's + # own requested_backend policy filter are not fallbacks. + capability_rejections = 0 + + for backend in candidates: + capability = self._logprob_capabilities.get(backend) + if capability is None: + rejected.append(f"{backend.name}: no LogprobBackendCapability declared") + capability_rejections += 1 + continue + capability_incompat = list(capability.incompatibilities(contract)) + policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) + reasons = capability_incompat + ([policy_mismatch] if policy_mismatch else []) + if reasons: + rejected.append(f"{backend.name}: " + "; ".join(reasons)) + if capability_incompat: + capability_rejections += 1 + continue + + op = self._get_or_create_backend(backend) + if op is None: + rejected.append(f"{backend.name}: backend could not be loaded or instantiated") + capability_rejections += 1 + continue + + provenance = { + "requested_backend": requested_backend, + "actual_backend": capability.backend_id, + "backend_enum": backend.name, + "platform": platform, + "fallback": capability_rejections > 0, + "prior_rejections": list(rejected), + "contract": contract.to_dict(), + "capability": capability.to_dict(), + } + return LogprobDispatchResult( + op=op, + capability=capability, + provenance=provenance, + ) + + details = " | ".join(rejected) if rejected else "no candidates registered" + requested = contract.to_dict() + raise RuntimeError( + "No logprob backend supports the requested WS2 contract on " + f"{platform}: role={requested['role']}, dtype={requested['dtype']}, " + f"TP={contract.sharding.tp_world_size}, CP={contract.sharding.cp_world_size}, " + f"padded_vocab={contract.sharding.padded_vocab_size}, " + f"real_vocab={contract.sharding.real_vocab_size}. Rejections: {details}" + ) + + @staticmethod + def _logprob_policy_mismatch( + requested_backend: str, + capability: LogprobBackendCapability, + ) -> str | None: + policy = requested_backend.lower() + if policy == "auto": + return None + if policy in {"production", "reference", "deterministic"}: + if capability.implementation_kind == policy: + return None + return ( + f"implementation_kind={capability.implementation_kind} does not satisfy " + f"requested_backend={policy}" + ) + if capability.backend_id == requested_backend: + return None + return ( + f"backend_id={capability.backend_id} does not match " + f"requested_backend={requested_backend}" + ) + + def _platform(self) -> str: + if device_ctx.is_rocm: + return "rocm" + if device_ctx.device_type == "cuda": + return "cuda" + return "cpu" + + def _get_or_create_backend(self, backend: OpBackend) -> Any | None: + if backend.name in self._instance_cache: + return self._instance_cache[backend.name] + if backend.name in self._failed_backends: + return None + + op_class = self._load_backend(backend) + if op_class is None: + self._failed_backends.add(backend.name) + return None + try: + op = op_class() + except Exception as exc: + logger.error(f"Failed to instantiate {backend.name}: {exc}") + self._failed_backends.add(backend.name) + return None + self._instance_cache[backend.name] = op + return op + def _platform_for_device(self, device: torch.device | str | None) -> str: if device is None: if device_ctx.is_rocm: diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py new file mode 100644 index 00000000..d18083f4 --- /dev/null +++ b/tests/test_logprob_contract.py @@ -0,0 +1,402 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2026 RL-Kernel Contributors + +"""WS2 TP-aware logprob contract and contract-aware dispatch tests (issue #241).""" + +from __future__ import annotations + +import json +from dataclasses import replace + +import pytest + +from rl_engine.kernels.logprob_contract import ( + LogprobBackendCapability, + LogprobContract, + LogprobContractError, + LogprobDType, + LogprobRole, + MaskSpec, + ReductionSpec, + ShardingSpec, +) +from rl_engine.kernels.registry import KernelRegistry, OpBackend + +QWEN3_REAL_VOCAB = 151936 +QWEN3_PADDED_VOCAB = 152064 + + +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)) + + +def _sharding( + *, + tp_rank: int = 0, + tp_world_size: int = 2, + cp_rank: int = 0, + cp_world_size: int = 2, + real_vocab_size: int = QWEN3_REAL_VOCAB, + padded_vocab_size: int = QWEN3_PADDED_VOCAB, + vocab_shard_bounds: tuple[tuple[int, int], ...] | None = None, +) -> ShardingSpec: + return ShardingSpec( + tp_rank=tp_rank, + tp_world_size=tp_world_size, + vocab_shard_bounds=( + vocab_shard_bounds + if vocab_shard_bounds is not None + else _even_bounds(padded_vocab_size, tp_world_size) + ), + real_vocab_size=real_vocab_size, + padded_vocab_size=padded_vocab_size, + cp_rank=cp_rank, + cp_world_size=cp_world_size, + ) + + +def _mask( + *, + num_tokens: int = 8, + active_mask: tuple[bool, ...] | None = None, + ignore_index: int = -100, +) -> MaskSpec: + return MaskSpec( + num_tokens=num_tokens, + active_mask=( + active_mask + if active_mask is not None + else (False, False, True, True, True, True, True, False) + ), + ignore_index=ignore_index, + ) + + +def _contract( + *, + role: str = "train", + dtype: str = "bf16", + mask: MaskSpec | None = None, + sharding: ShardingSpec | None = None, + reduction: ReductionSpec | None = None, +) -> LogprobContract: + return LogprobContract( + role=role, + dtype=dtype, + mask=mask if mask is not None else _mask(), + sharding=sharding if sharding is not None else _sharding(), + reduction=reduction if reduction is not None else ReductionSpec(), + ) + + +def _declared_tp_backend() -> LogprobBackendCapability: + return LogprobBackendCapability( + backend_id="test-deterministic-tp-logprob", + roles=frozenset({LogprobRole.TRAIN, LogprobRole.INFER}), + dtypes=frozenset({LogprobDType.BF16}), + tp_world_sizes=(1, 2, 4), + cp_world_sizes=None, + supports_vocab_padding=True, + supports_inactive_tokens=True, + exports_vocab_lse=True, + deterministic_tp_merge=True, + implementation_kind="deterministic", + ) + + +def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): + contract = _contract() + + assert contract.sharding.tp_world_size == 2 + assert contract.sharding.cp_world_size == 2 + assert contract.sharding.local_vocab_start == 0 + assert contract.sharding.local_vocab_end == QWEN3_PADDED_VOCAB // 2 + assert contract.sharding.local_vocab_size == QWEN3_PADDED_VOCAB // 2 + assert contract.mask.active_token_count == 5 + assert contract.reduction.acc_dtype is LogprobDType.FP32 + assert contract.to_dict()["reduction"] == { + "merge": "max_sumexp", + "merge_axis": "tp_vocab", + "acc_dtype": "fp32", + "order": "global_vocab_shard_index", + "transport": "all_gather", + "downcast_at": "final_write", + "engine": "in_op_reference", + "cp_is_merge_axis": False, + } + json.dumps(contract.to_dict()) + + +@pytest.mark.parametrize("tp_world_size", [1, 2, 4]) +def test_pr4_sweep_tp_degrees_are_representable(tp_world_size): + sharding = _sharding(tp_world_size=tp_world_size, cp_world_size=1) + + assert len(sharding.vocab_shard_bounds) == tp_world_size + assert sharding.vocab_shard_bounds[-1][1] == QWEN3_PADDED_VOCAB + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == tp_world_size - 1 + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("tp_rank", 2, "tp_rank=2"), + ("cp_rank", 2, "cp_rank=2"), + ("real_vocab_size", 0, "positive integer"), + ("padded_vocab_size", QWEN3_REAL_VOCAB - 1, "must not be smaller"), + ], +) +def test_invalid_rank_and_vocab_metadata_fail_loudly(field, value, message): + values = { + "tp_rank": 0, + "tp_world_size": 2, + "cp_rank": 0, + "cp_world_size": 2, + "real_vocab_size": QWEN3_REAL_VOCAB, + "padded_vocab_size": QWEN3_PADDED_VOCAB, + "vocab_shard_bounds": _even_bounds(QWEN3_PADDED_VOCAB, 2), + } + values[field] = value + + with pytest.raises(LogprobContractError, match=message): + ShardingSpec(**values) + + +@pytest.mark.parametrize( + ("bounds", "message"), + [ + ((), "one \\(start, end\\) pair per TP rank"), + (((0, 76032),), "one \\(start, end\\) pair per TP rank"), + (((0, 76032), (76032, 76032)), "end > start"), + (((0, 76000), (76032, 152064)), "contiguous"), + (((0, 76064), (76032, 152064)), "contiguous"), + (((0, 76032), (76032, 152000)), "cover padded_vocab_size exactly"), + ], +) +def test_incomplete_or_overlapping_vocab_shard_bounds_fail_loudly(bounds, message): + with pytest.raises(LogprobContractError, match=message): + _sharding(vocab_shard_bounds=bounds) + + +def test_owner_rank_is_unique_and_rejects_out_of_real_vocab_targets(): + sharding = _sharding() + + assert sharding.owner_rank(0) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2 - 1) == 0 + assert sharding.owner_rank(QWEN3_PADDED_VOCAB // 2) == 1 + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 1 + + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(-1) + with pytest.raises(LogprobContractError, match="outside the real vocabulary"): + sharding.owner_rank(QWEN3_REAL_VOCAB) + + +def test_active_token_mask_metadata_is_validated(): + with pytest.raises(LogprobContractError, match="one entry per token"): + _mask(num_tokens=4) + + with pytest.raises(LogprobContractError, match="must be a bool"): + MaskSpec(num_tokens=2, active_mask=(True, 1)) + + all_inactive = _mask(num_tokens=3, active_mask=(False, False, False)) + assert all_inactive.active_token_count == 0 + + +def test_reduction_requires_fp32_accumulation_and_known_semantics(): + with pytest.raises(LogprobContractError, match="must be fp32"): + ReductionSpec(acc_dtype="bf16") + + with pytest.raises(LogprobContractError, match="merge must be one of"): + ReductionSpec(merge="lse_average") + + with pytest.raises(LogprobContractError, match="transport must be one of"): + ReductionSpec(transport="all_reduce") + + +def test_contract_component_types_and_lse_export_are_enforced(): + with pytest.raises(LogprobContractError, match="mask must be a MaskSpec"): + LogprobContract( + role="train", + dtype="bf16", + mask=None, + sharding=_sharding(), + reduction=ReductionSpec(), + ) + + with pytest.raises(LogprobContractError, match="export_lse must be True"): + replace(_contract(), export_lse=False) + + +def test_ignore_index_must_not_collide_with_the_real_vocabulary(): + with pytest.raises(LogprobContractError, match="must not collide"): + _contract(mask=_mask(ignore_index=5)) + + padding_column = QWEN3_REAL_VOCAB + 1 + contract = _contract(mask=_mask(ignore_index=padding_column)) + assert contract.mask.ignore_index == padding_column + + +def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): + registry = KernelRegistry() + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(_contract()) + + message = str(exc_info.value) + assert "TP=2 is unsupported" in message + assert "vocab-domain LSE export is unsupported" in message + assert "deterministic TP (max, sumexp) merge is unsupported" in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_current_ws1_backend_rejects_padded_vocab_even_at_tp1(): + registry = KernelRegistry() + contract = _contract(sharding=_sharding(tp_world_size=1, cp_world_size=1)) + + with pytest.raises(RuntimeError) as exc_info: + registry.get_logprob_op(contract) + + message = str(exc_info.value) + assert "TP=1 is unsupported" not in message + assert "padded-vs-real vocab masking is unsupported" in message + + +def test_undeclared_backend_capability_is_never_selected(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_NATIVE] + + with pytest.raises(RuntimeError, match="no LogprobBackendCapability declared"): + registry.get_logprob_op(_contract()) + + +def test_declared_compatible_backend_resolves_and_records_provenance(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract(), requested_backend="deterministic") + + assert result.op is not None + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + assert result.provenance["fallback"] is False + assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 + assert result.provenance["contract"]["sharding"]["real_vocab_size"] == QWEN3_REAL_VOCAB + assert result.provenance["contract"]["reduction"]["cp_is_merge_axis"] is False + json.dumps(result.provenance) + + +def test_requested_stable_backend_id_is_enforced(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): + registry.get_logprob_op(_contract(), requested_backend="another-backend") + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" + + +def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): + capability = _declared_tp_backend() + cp2_contract = _contract(sharding=_sharding(cp_world_size=2, cp_rank=1)) + + assert capability.incompatibilities(cp2_contract) == () + + cp_restricted = replace(capability, cp_world_sizes=(1,)) + assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) + + +def test_inactive_tokens_require_declared_backend_support(): + capability = replace(_declared_tp_backend(), supports_inactive_tokens=False) + contract = _contract() + + assert "inactive-token (ignore_index) masking is unsupported" in ( + capability.incompatibilities(contract) + ) + + fully_active = _contract(mask=_mask(num_tokens=3, active_mask=(True, True, True))) + assert capability.incompatibilities(fully_active) == () + + +def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): + with pytest.raises(LogprobContractError, match="reserved dispatch policy keyword"): + replace(_declared_tp_backend(), backend_id="Deterministic") + + +def test_default_auto_policy_resolves_any_compatible_implementation_kind(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), implementation_kind="reference" + ) + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["requested_backend"] == "auto" + assert result.capability.implementation_kind == "reference" + + +def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract(), requested_backend="DETERMINISTIC") + assert result.capability.backend_id == "test-deterministic-tp-logprob" + + with pytest.raises(RuntimeError, match="does not match requested_backend"): + registry.get_logprob_op(_contract(), requested_backend="Test-Deterministic-TP-Logprob") + + +def test_policy_only_skips_are_not_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ] + registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), backend_id="other-compatible-backend" + ) + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_capability_rejections_are_reported_as_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [ + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + ] + registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( + _declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,) + ) + registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + + result = registry.get_logprob_op(_contract()) + + assert result.provenance["fallback"] is True + assert "TP=2 is unsupported" in result.provenance["prior_rejections"][0] + + +def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform].insert(0, OpBackend.PYTORCH_NATIVE) + + legacy = registry._priority_map[platform]["batch_invariant_logp"] + assert OpBackend.PYTORCH_NATIVE not in legacy From cdc11ba17621ef52e94f90fec8a477d0b9990075 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 22:50:48 +0800 Subject: [PATCH 2/6] fix(ws2): address CodeRabbit review on logprob contract PR MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 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 #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 --- docs/design/ws2-tp-logprob-contract.md | 16 +++-- rl_engine/kernels/logprob_contract.py | 8 ++- rl_engine/kernels/registry.py | 54 ++++++++++------ tests/test_logprob_contract.py | 88 +++++++++++++++++++------- 4 files changed, 116 insertions(+), 50 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index a392393b..3b1cf2fd 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -127,8 +127,15 @@ selected_logp = target_logit - LSE The selected target logit comes from a masked single-owner gather: exactly one rank holds each active token's target column. Downcast happens only at the final write. Because the -merge order is fixed by shard index, TP=2 is bitwise-equal to TP=1 by construction; averaging -per-rank logsumexp values or letting a collective reduce numerically is not conformant. +merge order is fixed by shard index, the result is deterministic and reproducible at every +TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 +acceptance target) requires one further condition: the local per-shard reduction must use a +TP-degree-independent tile decomposition, so that the same partial sums are formed in the +same order regardless of how the vocabulary is sharded. Providing that decomposition is an +obligation of the deterministic reference implementation; a backend without it is still +deterministic per degree, and its cross-degree drift is judged against the #108 tolerance +table instead. Averaging per-rank logsumexp values or letting a collective reduce +numerically is not conformant in either case. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift reports follow the #116 format. This contract does not introduce another tolerance table. @@ -170,8 +177,9 @@ The current WS1 batch-invariant logp implementations are single-shard (TP=1) ref they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference -becomes selectable by registering a capability that truthfully declares those features; no -controller branch or silent fallback is required. +becomes selectable through `KernelRegistry.register_logprob_backend(backend, capability)` +by declaring a capability that truthfully describes those features; no controller branch or +silent fallback is required. Successful dispatch provenance records: diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 8f2667e1..815d232f 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -342,10 +342,11 @@ def to_dict(self) -> dict[str, Any]: "engine": self.reduction.engine.value, "cp_is_merge_axis": False, } + # The per-token mask is deliberately summarized: provenance exists for + # logging/serialization and the raw mask would dominate its size. 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, } return { @@ -382,6 +383,7 @@ def __post_init__(self) -> None: raise LogprobContractError( f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" ) + object.__setattr__(self, "backend_id", self.backend_id.strip()) roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) if not roles or not dtypes: @@ -482,17 +484,17 @@ class LogprobDispatchResult: __all__ = [ + "RESERVED_DISPATCH_POLICIES", "DowncastPoint", "LogprobBackendCapability", "LogprobContract", "LogprobContractError", - "LogprobDispatchResult", "LogprobDType", + "LogprobDispatchResult", "LogprobMerge", "LogprobRole", "MaskSpec", "MergeAxis", - "RESERVED_DISPATCH_POLICIES", "ReductionEngine", "ReductionOrder", "ReductionSpec", diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 35ae951b..d9b741bf 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -433,25 +433,41 @@ def get_op(self, op_type: str, device: torch.device | str | None = None) -> Any: candidates = self._priority_map.get(platform, {}).get(op_type, [OpBackend.PYTORCH_NATIVE]) for backend in candidates: - if backend.name in self._instance_cache: - return self._instance_cache[backend.name] + op_instance = self._get_or_create_backend(backend) + if op_instance is not None: + return op_instance - if backend.name in self._failed_backends: - continue + raise RuntimeError(f"No functional backend found for {op_type} on {platform}") - 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) + def register_logprob_backend( + self, + backend: OpBackend, + capability: LogprobBackendCapability, + *, + platform: Optional[str] = None, + prepend: bool = False, + ) -> None: + """Register (or replace) a backend for WS2 contract-aware logprob dispatch. + + This is the supported seam for making a new backend selectable by + ``get_logprob_op`` (e.g. the deterministic vocab-parallel TP reference + from issue #241 PR 3) without touching the legacy ``get_op`` priority + lists. Registering the same backend again replaces its capability + without duplicating the candidate entry. + """ - raise RuntimeError(f"No functional backend found for {op_type} on {platform}") + if not isinstance(backend, OpBackend): + raise LogprobContractError("backend must be an OpBackend") + if not isinstance(capability, LogprobBackendCapability): + raise LogprobContractError("capability must be a LogprobBackendCapability") + resolved_platform = platform if platform is not None else self._platform() + candidates = self._logprob_candidates.setdefault(resolved_platform, []) + self._logprob_capabilities[backend] = capability + if backend not in candidates: + if prepend: + candidates.insert(0, backend) + else: + candidates.append(backend) def get_logprob_op( self, @@ -556,11 +572,7 @@ def _logprob_policy_mismatch( ) 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) def _get_or_create_backend(self, backend: OpBackend) -> Any | None: if backend.name in self._instance_cache: diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index d18083f4..8a7a1ec5 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -28,7 +28,10 @@ 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) + ) def _sharding( @@ -274,8 +277,10 @@ def test_undeclared_backend_capability_is_never_selected(): def test_declared_compatible_backend_resolves_and_records_provenance(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) result = registry.get_logprob_op(_contract(), requested_backend="deterministic") @@ -293,8 +298,10 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): def test_requested_stable_backend_id_is_enforced(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) with pytest.raises(RuntimeError, match="does not match requested_backend=another-backend"): registry.get_logprob_op(_contract(), requested_backend="another-backend") @@ -333,9 +340,11 @@ def test_backend_id_must_not_shadow_a_reserved_policy_keyword(): def test_default_auto_policy_resolves_any_compatible_implementation_kind(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), implementation_kind="reference" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), implementation_kind="reference"), + platform=platform, ) result = registry.get_logprob_op(_contract()) @@ -347,8 +356,10 @@ def test_default_auto_policy_resolves_any_compatible_implementation_kind(): def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) result = registry.get_logprob_op(_contract(), requested_backend="DETERMINISTIC") assert result.capability.backend_id == "test-deterministic-tp-logprob" @@ -360,14 +371,15 @@ def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): def test_policy_only_skips_are_not_reported_as_fallback(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [ + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( OpBackend.TRITON_BATCH_INVARIANT_LOGP, - OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, - ] - registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), backend_id="other-compatible-backend" + replace(_declared_tp_backend(), backend_id="other-compatible-backend"), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") @@ -378,14 +390,15 @@ def test_policy_only_skips_are_not_reported_as_fallback(): def test_capability_rejections_are_reported_as_fallback(): registry = KernelRegistry() platform = registry._platform() - registry._logprob_candidates[platform] = [ + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( OpBackend.TRITON_BATCH_INVARIANT_LOGP, - OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, - ] - registry._logprob_capabilities[OpBackend.TRITON_BATCH_INVARIANT_LOGP] = replace( - _declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,) + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - registry._logprob_capabilities[OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] = _declared_tp_backend() result = registry.get_logprob_op(_contract()) @@ -400,3 +413,34 @@ def test_ws2_candidate_list_is_decoupled_from_the_legacy_priority_map(): legacy = registry._priority_map[platform]["batch_invariant_logp"] assert OpBackend.PYTORCH_NATIVE not in legacy + + legacy.insert(0, OpBackend.PYTORCH_GEMM) + assert OpBackend.PYTORCH_GEMM not in registry._logprob_candidates[platform] + + +def test_register_logprob_backend_is_the_public_registration_seam(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + capability = _declared_tp_backend() + + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, capability, platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(capability, backend_id="replacement-backend"), + platform=platform, + ) + + assert registry._logprob_candidates[platform] == [OpBackend.PYTORCH_BATCH_INVARIANT_LOGP] + result = registry.get_logprob_op(_contract()) + assert result.capability.backend_id == "replacement-backend" + + with pytest.raises(LogprobContractError, match="capability must be"): + registry.register_logprob_backend(OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, None) + + +def test_backend_id_whitespace_is_normalized_for_dispatch(): + capability = replace(_declared_tp_backend(), backend_id=" padded-id ") + assert capability.backend_id == "padded-id" From 6455715eb09cfd4189f53a3aceb0a0df342a4dd8 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 23:06:28 +0800 Subject: [PATCH 3/6] fix(ws2): address second CodeRabbit round on logprob contract - 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 --- docs/design/ws2-tp-logprob-contract.md | 24 +++++++++------ rl_engine/kernels/logprob_contract.py | 15 ++++++--- rl_engine/kernels/registry.py | 22 ++++++++++++-- tests/test_logprob_contract.py | 42 ++++++++++++++++++++++++++ 4 files changed, 87 insertions(+), 16 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index 3b1cf2fd..28685bb4 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -114,9 +114,11 @@ downcast_at: final_write engine: in_op_reference ``` -Every rank computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` in -fp32, the partials travel by all-gather (collectives are transport only, never a numerical -reduction), and every rank merges in fixed global vocab-shard index order: +Every rank first masks every local column whose global id lies in +`[real_vocab_size, padded_vocab_size)` to `-inf`, so padding never contributes to the +logsumexp, then computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` +in fp32. The partials travel by all-gather (collectives are transport only, never a +numerical reduction), and every rank merges in fixed global vocab-shard index order: ```text M = max_l(m_l) @@ -129,12 +131,16 @@ The selected target logit comes from a masked single-owner gather: exactly one r each active token's target column. Downcast happens only at the final write. Because the merge order is fixed by shard index, the result is deterministic and reproducible at every TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 -acceptance target) requires one further condition: the local per-shard reduction must use a -TP-degree-independent tile decomposition, so that the same partial sums are formed in the -same order regardless of how the vocabulary is sharded. Providing that decomposition is an -obligation of the deterministic reference implementation; a backend without it is still -deterministic per degree, and its cross-degree drift is judged against the #108 tolerance -table instead. Averaging per-rank logsumexp values or letting a collective reduce +acceptance target) requires one further condition: the entire reduction must follow a +global tile-level structure that is independent of TP partitioning — a fixed tile +decomposition of the vocabulary plus a fixed merge order and rescaling tree over those +tiles, identical at every TP degree, so that the TP degree only selects which rank computes +which tiles and never changes the floating-point grouping. A TP-degree-independent +decomposition inside each shard is not sufficient on its own, because shard boundaries +would still group the combines differently across degrees. Providing that global structure +is an obligation of the deterministic reference implementation; a backend without it is +still deterministic per degree, and its cross-degree drift is judged against the #108 +tolerance table instead. Averaging per-rank logsumexp values or letting a collective reduce numerically is not conformant in either case. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 815d232f..2cb22a25 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -30,6 +30,9 @@ # Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable # backend id must never shadow one of these, or it becomes unselectable by id. RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) +# Policies a backend can declare as its implementation kind; "auto" is a +# selection strategy, not an implementation kind. +IMPLEMENTATION_KINDS = RESERVED_DISPATCH_POLICIES - {"auto"} class LogprobContractError(ValueError): @@ -384,8 +387,11 @@ def __post_init__(self) -> None: f"backend_id={self.backend_id!r} shadows a reserved dispatch policy keyword" ) object.__setattr__(self, "backend_id", self.backend_id.strip()) - roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) - dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + try: + roles = frozenset(_enum_value(LogprobRole, value, "roles") for value in self.roles) + dtypes = frozenset(_enum_value(LogprobDType, value, "dtypes") for value in self.dtypes) + except TypeError as exc: + raise LogprobContractError("roles and dtypes must be iterables of enum values") from exc if not roles or not dtypes: raise LogprobContractError("backend roles and dtypes must not be empty") tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") @@ -398,9 +404,9 @@ def __post_init__(self) -> None: ): if not isinstance(getattr(self, flag_name), bool): raise LogprobContractError(f"{flag_name} must be a bool") - 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" + f"implementation_kind must be one of: {', '.join(sorted(IMPLEMENTATION_KINDS))}" ) object.__setattr__(self, "roles", roles) object.__setattr__(self, "dtypes", dtypes) @@ -484,6 +490,7 @@ class LogprobDispatchResult: __all__ = [ + "IMPLEMENTATION_KINDS", "RESERVED_DISPATCH_POLICIES", "DowncastPoint", "LogprobBackendCapability", diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index d9b741bf..0f4f2875 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -181,7 +181,7 @@ def __init__(self): # of silently selecting an incompatible fallback. common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) - self._logprob_capabilities = { + base_logprob_capabilities = { OpBackend.PYTORCH_BATCH_INVARIANT_LOGP: LogprobBackendCapability( backend_id="pytorch-batch-invariant-logp-ws1", roles=common_logprob_roles, @@ -343,6 +343,16 @@ def __init__(self): platform: list(ops.get("batch_invariant_logp", [])) for platform, ops in self._priority_map.items() } + # Capabilities are scoped per platform: the same backend enum may + # truthfully declare different support on cuda vs rocm vs cpu. + self._logprob_capabilities: Dict[str, Dict[OpBackend, LogprobBackendCapability]] = { + platform: { + backend: base_logprob_capabilities[backend] + for backend in candidates + if backend in base_logprob_capabilities + } + for platform, candidates in self._logprob_candidates.items() + } def _adjust_priority_from_env(self): rocm_attn_backend = os.getenv("RL_KERNEL_ROCM_ATTN_BACKEND", "").strip().lower() @@ -461,8 +471,13 @@ def register_logprob_backend( if not isinstance(capability, LogprobBackendCapability): raise LogprobContractError("capability must be a LogprobBackendCapability") resolved_platform = platform if platform is not None else self._platform() + if resolved_platform not in self._priority_map: + raise LogprobContractError( + f"unsupported platform {resolved_platform!r}; expected one of " + f"{sorted(self._priority_map)}" + ) candidates = self._logprob_candidates.setdefault(resolved_platform, []) - self._logprob_capabilities[backend] = capability + self._logprob_capabilities.setdefault(resolved_platform, {})[backend] = capability if backend not in candidates: if prepend: candidates.insert(0, backend) @@ -502,8 +517,9 @@ def get_logprob_op( # own requested_backend policy filter are not fallbacks. capability_rejections = 0 + platform_capabilities = self._logprob_capabilities.get(platform, {}) for backend in candidates: - capability = self._logprob_capabilities.get(backend) + capability = platform_capabilities.get(backend) if capability is None: rejected.append(f"{backend.name}: no LogprobBackendCapability declared") capability_rejections += 1 diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 8a7a1ec5..670225f0 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -444,3 +444,45 @@ def test_register_logprob_backend_is_the_public_registration_seam(): def test_backend_id_whitespace_is_normalized_for_dispatch(): capability = replace(_declared_tp_backend(), backend_id=" padded-id ") assert capability.backend_id == "padded-id" + + +def test_capabilities_are_scoped_per_platform(): + registry = KernelRegistry() + platform = registry._platform() + other = "rocm" if platform != "rocm" else "cpu" + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="other-platform-backend"), + platform=other, + ) + + result = registry.get_logprob_op(_contract()) + + assert result.capability.backend_id == "test-deterministic-tp-logprob" + assert ( + registry._logprob_capabilities[other][OpBackend.PYTORCH_BATCH_INVARIANT_LOGP].backend_id + == "other-platform-backend" + ) + + +def test_register_logprob_backend_rejects_unknown_platform(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="unsupported platform"): + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, + _declared_tp_backend(), + platform="cuda-typo", + ) + + +def test_non_iterable_roles_and_dtypes_raise_contract_errors(): + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), roles=None) + + with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): + replace(_declared_tp_backend(), dtypes=42) From 3b4eaef3dd7ccb3645653e79166b03d0492355d1 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Sun, 2 Aug 2026 23:48:14 +0800 Subject: [PATCH 4/6] feat(ws2): make determinism scope and invocation surface part of the 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 #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 --- docs/design/ws2-tp-logprob-contract.md | 68 +++++++--- rl_engine/kernels/logprob_contract.py | 175 ++++++++++++++++++++++--- rl_engine/kernels/registry.py | 42 +++--- tests/test_logprob_contract.py | 116 ++++++++++++++-- 4 files changed, 343 insertions(+), 58 deletions(-) diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md index 28685bb4..cf9179b8 100644 --- a/docs/design/ws2-tp-logprob-contract.md +++ b/docs/design/ws2-tp-logprob-contract.md @@ -36,12 +36,17 @@ for provenance and must never widen the merge. `rl_engine.kernels.logprob_contract` defines: -- `LogprobContract`: role, logits dtype, mask, sharding, reduction, and LSE export; +- `LogprobContract`: role, logits dtype, mask, sharding, reduction, output surface, and + LSE export, plus a rank-independent `cross_rank_fingerprint()`; - `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank metadata, and target-token ownership; - `MaskSpec`: active-token mask and ignore index; -- `ReductionSpec`: fixed `(max, sumexp)` merge semantics; -- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports. +- `ReductionSpec`: fixed `(max, sumexp)` merge semantics and the requested determinism + scope; +- `LogprobOutputSpec`: the output surface — fp32 selected logprob and fp32 vocab-domain + LSE, replicated across the TP group; +- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports, + including its mask modes and determinism scopes. Construction performs validation immediately. A structurally valid contract means that the request is complete and internally consistent; it does not mean that an installed backend can @@ -140,8 +145,21 @@ decomposition inside each shard is not sufficient on its own, because shard boun would still group the combines differently across degrees. Providing that global structure is an obligation of the deterministic reference implementation; a backend without it is still deterministic per degree, and its cross-degree drift is judged against the #108 -tolerance table instead. Averaging per-rank logsumexp values or letting a collective reduce -numerically is not conformant in either case. +tolerance table instead. The contract expresses this distinction as +`ReductionSpec.determinism_scope`: `cross_tp_bitwise` (the #241 target and the default) +versus `fixed_topology`. A backend declares the scopes it honors in +`LogprobBackendCapability.determinism_scopes`, and dispatch rejects a backend that cannot +honor the requested scope — prose obligations are not enough; the guarantee is part of the +typed contract. + +A shard may lie entirely inside the padded region, and a row's local columns may all be +`-inf` after masking. The identity partial for these cases is defined as +`(m_l, s_l) = (-inf, 0)`: a partial with `s_l = 0` contributes nothing to the merge +regardless of its `m_l`, and implementations must use this identity directly rather than +evaluating `exp(-inf - (-inf))`, which would poison the merge with NaN. + +Averaging per-rank logsumexp values or letting a collective reduce numerically is never +conformant, at either determinism scope. The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift reports follow the #116 format. This contract does not introduce another tolerance table. @@ -164,16 +182,20 @@ provenance = result.provenance ``` Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, -TP/CP degree, padded-vs-real vocab masking, inactive-token support, vocab-domain LSE export, -and deterministic TP merge. An undeclared or incompatible backend is skipped with an -explicit rejection reason; there is no silent fallback. +TP/CP degree, padded-vs-real vocab masking, explicit active-mask support, vocab-domain LSE +export, and the requested determinism scope. An undeclared or incompatible backend is +skipped with an explicit rejection reason; there is no silent fallback. `requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | -`reference` | `deterministic`; default `auto`) or an exact, case-sensitive stable backend -id. Strictness comes from the contract's capability checks, not from the policy string. A -backend id may never shadow a policy keyword; capability construction rejects that. The -provenance `fallback` flag reports only capability or load rejections of otherwise-eligible -candidates — skips caused purely by the caller's own policy filter are not fallbacks. +`reference`; default `auto`) or an exact, case-sensitive stable backend id. The keywords +select an implementation tier; determinism is not a tier — it is requested through +`ReductionSpec.determinism_scope`, so `requested_backend="deterministic"` raises a loud +error instead of silently matching nothing. Strictness comes from the contract's +capability checks, not from the policy string. A backend id may never shadow a reserved +keyword; capability construction rejects that. The provenance `fallback` flag reports only +capability or load rejections of policy-eligible candidates — a candidate excluded by the +caller's own policy never counts as a fallback, even if it would also have failed +capability checks. WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy `batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does @@ -192,10 +214,26 @@ Successful dispatch provenance records: - requested and actual backend ids; - platform and fallback status; - prior candidate rejection reasons; -- the complete requested contract, including shard bounds, padded and real vocab sizes, - merge semantics, and the explicit `cp_is_merge_axis: false` declaration; +- the complete dispatch-relevant contract, including shard bounds, padded and real vocab + sizes, merge and output semantics, the explicit `cp_is_merge_axis: false` declaration, + and the active-mask digest (`active_mask_sha256`) — the mask's identity without its + per-token payload; - the selected backend capability descriptor. +### Distributed dispatch safety + +`get_logprob_op` resolves locally on each rank, so `requested_backend="auto"` is not +distributed-safe on its own: a load failure on one rank can resolve a different backend +than its peers, which for a collective-bearing implementation means divergent numerical +schedules or a deadlock. For `tp_world_size > 1` a caller must either request an exact +backend id or run a preflight agreement before any collective: all-gather the resolved +backend id together with `LogprobContract.cross_rank_fingerprint()` — a rank-independent +hash covering the shard-bounds table, vocab sizes, reduction/output semantics, and the +active-mask digest, excluding rank-local fields — and abort on any mismatch. Implementing +this preflight is an obligation of the #241 PR 3/PR 4 work; the backend invocation +protocol (how the contract and mask reach the implementation) is likewise defined there, +against this contract. + ## Validation Contract and dispatch behavior are covered by: diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 2cb22a25..6941c78e 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -21,6 +21,8 @@ from __future__ import annotations +import hashlib +import json from dataclasses import dataclass, field from enum import Enum from typing import Any, TypeVar @@ -29,10 +31,15 @@ # Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable # backend id must never shadow one of these, or it becomes unselectable by id. +# "deterministic" stays reserved even though it is no longer a policy: +# determinism is expressed through DeterminismScope, and requesting it as a +# policy is a loud error rather than a silent id mismatch. RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) -# Policies a backend can declare as its implementation kind; "auto" is a -# selection strategy, not an implementation kind. -IMPLEMENTATION_KINDS = RESERVED_DISPATCH_POLICIES - {"auto"} +# Implementation tiers a backend can declare. Determinism is deliberately a +# separate axis (DeterminismScope): a backend can be a deterministic reference, +# a deterministic production implementation, or a non-deterministic production +# implementation. +IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) class LogprobContractError(ValueError): @@ -85,6 +92,40 @@ class ReductionEngine(str, Enum): IN_OP_REFERENCE = "in_op_reference" +class DeterminismScope(str, Enum): + """Strength of the reduction's determinism guarantee. + + ``fixed_topology``: bitwise-reproducible for one fixed TP degree; results + at different TP degrees are compared against the #108 tolerance table. + + ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This + requires the entire reduction to follow a global tile-level structure that + is independent of TP partitioning (see the design doc); fixed shard-order + merging alone is not sufficient. + """ + + FIXED_TOPOLOGY = "fixed_topology" + CROSS_TP_BITWISE = "cross_tp_bitwise" + + +class MaskMode(str, Enum): + """How a backend consumes inactive-token information. + + ``explicit_active_mask``: the backend honors an arbitrary active-token + mask. ``ignore_index``: the backend only recognizes inactive tokens whose + target id equals ``ignore_index``. The contract permits inactive targets + that do NOT hold ``ignore_index``, so an ignore-index-only backend cannot + serve a contract with inactive tokens. + """ + + EXPLICIT_ACTIVE_MASK = "explicit_active_mask" + IGNORE_INDEX = "ignore_index" + + +class TPPlacement(str, Enum): + REPLICATED = "replicated" + + def _enum_value(enum_type: type[_EnumT], value: Any, field: str) -> _EnumT: try: return enum_type(value) @@ -233,6 +274,7 @@ class MaskSpec: active_mask: tuple[bool, ...] ignore_index: int = -100 _active_token_count: int = field(init=False, repr=False, compare=False) + _active_mask_sha256: str = field(init=False, repr=False, compare=False) def __post_init__(self) -> None: num_tokens = _positive_int(self.num_tokens, "num_tokens") @@ -251,11 +293,19 @@ def __post_init__(self) -> None: ) object.__setattr__(self, "active_mask", active_mask) object.__setattr__(self, "_active_token_count", sum(active_mask)) + object.__setattr__( + self, "_active_mask_sha256", hashlib.sha256(bytes(active_mask)).hexdigest() + ) @property def active_token_count(self) -> int: return self._active_token_count + @property + def active_mask_sha256(self) -> str: + """Compact mask identity for provenance and cross-rank agreement.""" + return self._active_mask_sha256 + @dataclass(frozen=True) class ReductionSpec: @@ -268,8 +318,14 @@ class ReductionSpec: transport: ReductionTransport = ReductionTransport.ALL_GATHER downcast_at: DowncastPoint = DowncastPoint.FINAL_WRITE engine: ReductionEngine = ReductionEngine.IN_OP_REFERENCE + determinism_scope: DeterminismScope = DeterminismScope.CROSS_TP_BITWISE def __post_init__(self) -> None: + object.__setattr__( + self, + "determinism_scope", + _enum_value(DeterminismScope, self.determinism_scope, "determinism_scope"), + ) object.__setattr__(self, "merge", _enum_value(LogprobMerge, self.merge, "merge")) object.__setattr__( self, "merge_axis", _enum_value(MergeAxis, self.merge_axis, "merge_axis") @@ -291,6 +347,39 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class LogprobOutputSpec: + """Output surface every conforming backend must produce. + + Selected logprob and vocab-domain LSE are fp32 and replicated across the + TP group; the ``downcast_at: final_write`` rule applies to any consumer + downcast after these outputs, never inside the reduction. + """ + + selected_logp_dtype: LogprobDType = LogprobDType.FP32 + lse_dtype: LogprobDType = LogprobDType.FP32 + tp_placement: TPPlacement = TPPlacement.REPLICATED + + def __post_init__(self) -> None: + object.__setattr__( + self, + "selected_logp_dtype", + _enum_value(LogprobDType, self.selected_logp_dtype, "selected_logp_dtype"), + ) + object.__setattr__( + self, "lse_dtype", _enum_value(LogprobDType, self.lse_dtype, "lse_dtype") + ) + object.__setattr__( + self, "tp_placement", _enum_value(TPPlacement, self.tp_placement, "tp_placement") + ) + if self.selected_logp_dtype is not LogprobDType.FP32: + raise LogprobContractError( + f"selected logprob output must be fp32; got {self.selected_logp_dtype.value}" + ) + if self.lse_dtype is not LogprobDType.FP32: + raise LogprobContractError(f"vocab LSE output must be fp32; got {self.lse_dtype.value}") + + @dataclass(frozen=True) class LogprobContract: """Complete semantic request consumed by contract-aware dispatch.""" @@ -300,6 +389,7 @@ class LogprobContract: mask: MaskSpec sharding: ShardingSpec reduction: ReductionSpec + output: LogprobOutputSpec = field(default_factory=LogprobOutputSpec) export_lse: bool = True def __post_init__(self) -> None: @@ -311,6 +401,8 @@ def __post_init__(self) -> None: raise LogprobContractError("sharding must be a ShardingSpec") if not isinstance(self.reduction, ReductionSpec): raise LogprobContractError("reduction must be a ReductionSpec") + if not isinstance(self.output, LogprobOutputSpec): + raise LogprobContractError("output must be a LogprobOutputSpec") if not isinstance(self.export_lse, bool) or not self.export_lse: raise LogprobContractError( "export_lse must be True for the WS2 vocab-domain LSE drift contract" @@ -343,15 +435,24 @@ def to_dict(self) -> dict[str, Any]: "transport": self.reduction.transport.value, "downcast_at": self.reduction.downcast_at.value, "engine": self.reduction.engine.value, + "determinism_scope": self.reduction.determinism_scope.value, "cp_is_merge_axis": False, } # The per-token mask is deliberately summarized: provenance exists for - # logging/serialization and the raw mask would dominate its size. + # logging/serialization and the raw mask would dominate its size. The + # digest keeps the mask *identity* observable, so two masks with the + # same active count still produce distinguishable provenance. mask = { "num_tokens": self.mask.num_tokens, "active_token_count": self.mask.active_token_count, + "active_mask_sha256": self.mask.active_mask_sha256, "ignore_index": self.mask.ignore_index, } + output = { + "selected_logp_dtype": self.output.selected_logp_dtype.value, + "lse_dtype": self.output.lse_dtype.value, + "tp_placement": self.output.tp_placement.value, + } return { "semantic_operator": "selected_token_logprob", "role": self.role.value, @@ -361,7 +462,28 @@ def to_dict(self) -> dict[str, Any]: "mask": mask, "sharding": sharding, "reduction": reduction, + "output": output, + } + + 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() @dataclass(frozen=True) @@ -374,9 +496,9 @@ class LogprobBackendCapability: tp_world_sizes: tuple[int, ...] | None = None cp_world_sizes: tuple[int, ...] | None = None supports_vocab_padding: bool = False - supports_inactive_tokens: bool = False + mask_modes: frozenset[MaskMode] = frozenset() exports_vocab_lse: bool = False - deterministic_tp_merge: bool = False + determinism_scopes: frozenset[DeterminismScope] = frozenset() implementation_kind: str = "production" def __post_init__(self) -> None: @@ -396,12 +518,19 @@ def __post_init__(self) -> None: raise LogprobContractError("backend roles and dtypes must not be empty") tp_world_sizes = self._validated_world_sizes(self.tp_world_sizes, "tp_world_sizes") cp_world_sizes = self._validated_world_sizes(self.cp_world_sizes, "cp_world_sizes") - for flag_name in ( - "supports_vocab_padding", - "supports_inactive_tokens", - "exports_vocab_lse", - "deterministic_tp_merge", - ): + try: + mask_modes = frozenset( + _enum_value(MaskMode, value, "mask_modes") for value in self.mask_modes + ) + determinism_scopes = frozenset( + _enum_value(DeterminismScope, value, "determinism_scopes") + for value in self.determinism_scopes + ) + except TypeError as exc: + raise LogprobContractError( + "mask_modes and determinism_scopes must be iterables of enum values" + ) from exc + for flag_name in ("supports_vocab_padding", "exports_vocab_lse"): if not isinstance(getattr(self, flag_name), bool): raise LogprobContractError(f"{flag_name} must be a bool") if self.implementation_kind not in IMPLEMENTATION_KINDS: @@ -412,6 +541,8 @@ def __post_init__(self) -> None: object.__setattr__(self, "dtypes", dtypes) object.__setattr__(self, "tp_world_sizes", tp_world_sizes) object.__setattr__(self, "cp_world_sizes", cp_world_sizes) + object.__setattr__(self, "mask_modes", mask_modes) + object.__setattr__(self, "determinism_scopes", determinism_scopes) @staticmethod def _validated_world_sizes( @@ -453,13 +584,17 @@ def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: reasons.append("padded-vs-real vocab masking is unsupported") if ( contract.mask.active_token_count != contract.mask.num_tokens - and not self.supports_inactive_tokens + and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes ): - reasons.append("inactive-token (ignore_index) masking is unsupported") + # The contract does not require inactive targets to hold + # ignore_index, so ignore-index-only masking is insufficient. + reasons.append("explicit active-token masking is unsupported") if contract.export_lse and not self.exports_vocab_lse: reasons.append("vocab-domain LSE export is unsupported") - if tp_size > 1 and not self.deterministic_tp_merge: - reasons.append("deterministic TP (max, sumexp) merge is unsupported") + if contract.reduction.determinism_scope not in self.determinism_scopes: + reasons.append( + f"determinism_scope={contract.reduction.determinism_scope.value} is unsupported" + ) return tuple(reasons) def supports(self, contract: LogprobContract) -> bool: @@ -473,9 +608,9 @@ def to_dict(self) -> dict[str, Any]: "tp_world_sizes": list(self.tp_world_sizes) if self.tp_world_sizes else None, "cp_world_sizes": list(self.cp_world_sizes) if self.cp_world_sizes else None, "supports_vocab_padding": self.supports_vocab_padding, - "supports_inactive_tokens": self.supports_inactive_tokens, + "mask_modes": sorted(mode.value for mode in self.mask_modes), "exports_vocab_lse": self.exports_vocab_lse, - "deterministic_tp_merge": self.deterministic_tp_merge, + "determinism_scopes": sorted(scope.value for scope in self.determinism_scopes), "implementation_kind": self.implementation_kind, } @@ -492,6 +627,7 @@ class LogprobDispatchResult: __all__ = [ "IMPLEMENTATION_KINDS", "RESERVED_DISPATCH_POLICIES", + "DeterminismScope", "DowncastPoint", "LogprobBackendCapability", "LogprobContract", @@ -499,7 +635,9 @@ class LogprobDispatchResult: "LogprobDType", "LogprobDispatchResult", "LogprobMerge", + "LogprobOutputSpec", "LogprobRole", + "MaskMode", "MaskSpec", "MergeAxis", "ReductionEngine", @@ -507,4 +645,5 @@ class LogprobDispatchResult: "ReductionSpec", "ReductionTransport", "ShardingSpec", + "TPPlacement", ] diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index 0f4f2875..b17acf3e 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -9,12 +9,15 @@ import torch from rl_engine.kernels.logprob_contract import ( + IMPLEMENTATION_KINDS, + DeterminismScope, LogprobBackendCapability, LogprobContract, LogprobContractError, LogprobDispatchResult, LogprobDType, LogprobRole, + MaskMode, ) from rl_engine.platforms.device import device_ctx from rl_engine.utils.logger import logger @@ -188,9 +191,9 @@ def __init__(self): dtypes=common_logprob_dtypes, tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), implementation_kind="reference", ), OpBackend.TRITON_BATCH_INVARIANT_LOGP: LogprobBackendCapability( @@ -199,10 +202,10 @@ def __init__(self): dtypes=common_logprob_dtypes, tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, - implementation_kind="deterministic", + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", ), OpBackend.CUDA_BATCH_INVARIANT_LOGP_SM90: LogprobBackendCapability( backend_id="cuda-batch-invariant-logp-sm90-ws1", @@ -210,10 +213,10 @@ def __init__(self): dtypes=frozenset({LogprobDType.BF16, LogprobDType.FP32}), tp_world_sizes=(1,), supports_vocab_padding=False, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.IGNORE_INDEX}), exports_vocab_lse=False, - deterministic_tp_merge=False, - implementation_kind="deterministic", + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + implementation_kind="production", ), } @@ -508,6 +511,12 @@ def get_logprob_op( if not isinstance(requested_backend, str) or not requested_backend.strip(): raise LogprobContractError("requested_backend must be a non-empty string") requested_backend = requested_backend.strip() + if requested_backend.lower() == "deterministic": + raise LogprobContractError( + 'requested_backend="deterministic" is not a dispatch policy; request ' + "determinism through ReductionSpec.determinism_scope and match it against " + "backend determinism_scopes instead" + ) platform = self._platform() candidates = self._logprob_candidates.get(platform, []) @@ -524,13 +533,16 @@ def get_logprob_op( rejected.append(f"{backend.name}: no LogprobBackendCapability declared") capability_rejections += 1 continue - capability_incompat = list(capability.incompatibilities(contract)) policy_mismatch = self._logprob_policy_mismatch(requested_backend, capability) - reasons = capability_incompat + ([policy_mismatch] if policy_mismatch else []) - if reasons: - rejected.append(f"{backend.name}: " + "; ".join(reasons)) - if capability_incompat: - capability_rejections += 1 + if policy_mismatch is not None: + # Excluded by the caller's own policy: never a fallback, even + # if the candidate would also have failed capability checks. + rejected.append(f"{backend.name}: {policy_mismatch}") + continue + capability_incompat = list(capability.incompatibilities(contract)) + if capability_incompat: + rejected.append(f"{backend.name}: " + "; ".join(capability_incompat)) + capability_rejections += 1 continue op = self._get_or_create_backend(backend) @@ -573,7 +585,7 @@ def _logprob_policy_mismatch( policy = requested_backend.lower() if policy == "auto": return None - if policy in {"production", "reference", "deterministic"}: + if policy in IMPLEMENTATION_KINDS: if capability.implementation_kind == policy: return None return ( diff --git a/tests/test_logprob_contract.py b/tests/test_logprob_contract.py index 670225f0..40cea4e2 100644 --- a/tests/test_logprob_contract.py +++ b/tests/test_logprob_contract.py @@ -11,14 +11,18 @@ import pytest from rl_engine.kernels.logprob_contract import ( + DeterminismScope, LogprobBackendCapability, LogprobContract, LogprobContractError, LogprobDType, + LogprobOutputSpec, LogprobRole, + MaskMode, MaskSpec, ReductionSpec, ShardingSpec, + TPPlacement, ) from rl_engine.kernels.registry import KernelRegistry, OpBackend @@ -101,10 +105,12 @@ def _declared_tp_backend() -> LogprobBackendCapability: tp_world_sizes=(1, 2, 4), cp_world_sizes=None, supports_vocab_padding=True, - supports_inactive_tokens=True, + mask_modes=frozenset({MaskMode.EXPLICIT_ACTIVE_MASK, MaskMode.IGNORE_INDEX}), exports_vocab_lse=True, - deterministic_tp_merge=True, - implementation_kind="deterministic", + determinism_scopes=frozenset( + {DeterminismScope.CROSS_TP_BITWISE, DeterminismScope.FIXED_TOPOLOGY} + ), + implementation_kind="reference", ) @@ -126,6 +132,7 @@ def test_qwen3_tp2_bf16_contract_is_representable_and_serializable(): "transport": "all_gather", "downcast_at": "final_write", "engine": "in_op_reference", + "determinism_scope": "cross_tp_bitwise", "cp_is_merge_axis": False, } json.dumps(contract.to_dict()) @@ -249,7 +256,7 @@ def test_current_ws1_backend_rejects_strict_tp_contract_without_fallback(): message = str(exc_info.value) assert "TP=2 is unsupported" in message assert "vocab-domain LSE export is unsupported" in message - assert "deterministic TP (max, sumexp) merge is unsupported" in message + assert "determinism_scope=cross_tp_bitwise is unsupported" in message assert "padded-vs-real vocab masking is unsupported" in message @@ -282,11 +289,11 @@ def test_declared_compatible_backend_resolves_and_records_provenance(): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - result = registry.get_logprob_op(_contract(), requested_backend="deterministic") + result = registry.get_logprob_op(_contract(), requested_backend="reference") assert result.op is not None assert result.capability.backend_id == "test-deterministic-tp-logprob" - assert result.provenance["requested_backend"] == "deterministic" + assert result.provenance["requested_backend"] == "reference" assert result.provenance["actual_backend"] == "test-deterministic-tp-logprob" assert result.provenance["fallback"] is False assert result.provenance["contract"]["sharding"]["tp_world_size"] == 2 @@ -320,11 +327,11 @@ def test_cp_is_a_non_merge_axis_and_cp_agnostic_backends_accept_any_cp_degree(): assert cp_restricted.incompatibilities(cp2_contract) == ("CP=2 is unsupported",) -def test_inactive_tokens_require_declared_backend_support(): - capability = replace(_declared_tp_backend(), supports_inactive_tokens=False) +def test_inactive_tokens_require_explicit_active_mask_support(): + capability = replace(_declared_tp_backend(), mask_modes=frozenset({MaskMode.IGNORE_INDEX})) contract = _contract() - assert "inactive-token (ignore_index) masking is unsupported" in ( + assert "explicit active-token masking is unsupported" in ( capability.incompatibilities(contract) ) @@ -361,7 +368,7 @@ def test_policy_keywords_are_case_insensitive_but_backend_ids_are_exact(): OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform ) - result = registry.get_logprob_op(_contract(), requested_backend="DETERMINISTIC") + result = registry.get_logprob_op(_contract(), requested_backend="REFERENCE") assert result.capability.backend_id == "test-deterministic-tp-logprob" with pytest.raises(RuntimeError, match="does not match requested_backend"): @@ -486,3 +493,92 @@ def test_non_iterable_roles_and_dtypes_raise_contract_errors(): with pytest.raises(LogprobContractError, match="roles and dtypes must be iterables"): replace(_declared_tp_backend(), dtypes=42) + + +def test_requested_deterministic_policy_is_a_loud_error(): + registry = KernelRegistry() + + with pytest.raises(LogprobContractError, match="determinism_scope"): + registry.get_logprob_op(_contract(), requested_backend="deterministic") + + +def test_determinism_scope_is_part_of_the_typed_contract(): + fixed_only = replace( + _declared_tp_backend(), + determinism_scopes=frozenset({DeterminismScope.FIXED_TOPOLOGY}), + ) + + assert "determinism_scope=cross_tp_bitwise is unsupported" in ( + fixed_only.incompatibilities(_contract()) + ) + + relaxed = _contract(reduction=ReductionSpec(determinism_scope="fixed_topology")) + assert fixed_only.incompatibilities(relaxed) == () + + +def test_policy_filtered_candidates_never_count_toward_fallback(): + registry = KernelRegistry() + platform = registry._platform() + registry._logprob_candidates[platform] = [] + registry.register_logprob_backend( + OpBackend.TRITON_BATCH_INVARIANT_LOGP, + replace(_declared_tp_backend(), backend_id="tp1-only-backend", tp_world_sizes=(1,)), + platform=platform, + ) + registry.register_logprob_backend( + OpBackend.PYTORCH_BATCH_INVARIANT_LOGP, _declared_tp_backend(), platform=platform + ) + + result = registry.get_logprob_op(_contract(), requested_backend="test-deterministic-tp-logprob") + + assert result.provenance["fallback"] is False + assert len(result.provenance["prior_rejections"]) == 1 + + +def test_output_spec_is_pinned_to_fp32_replicated(): + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(selected_logp_dtype="bf16") + with pytest.raises(LogprobContractError, match="must be fp32"): + LogprobOutputSpec(lse_dtype="bf16") + + assert LogprobOutputSpec().tp_placement is TPPlacement.REPLICATED + assert _contract().to_dict()["output"] == { + "selected_logp_dtype": "fp32", + "lse_dtype": "fp32", + "tp_placement": "replicated", + } + + +def test_cross_rank_fingerprint_is_rank_independent_and_content_sensitive(): + rank0 = _contract(sharding=_sharding(tp_rank=0)) + rank1 = _contract(sharding=_sharding(tp_rank=1, cp_rank=1)) + + assert rank0.cross_rank_fingerprint() == rank1.cross_rank_fingerprint() + + different_mask = _contract( + mask=_mask(active_mask=(True, True, True, True, True, True, True, False)) + ) + assert rank0.cross_rank_fingerprint() != different_mask.cross_rank_fingerprint() + + +def test_provenance_records_the_active_mask_digest(): + provenance_mask = _contract().to_dict()["mask"] + + assert provenance_mask["active_mask_sha256"] == _mask().active_mask_sha256 + assert len(provenance_mask["active_mask_sha256"]) == 64 + + same_count_different_mask = _mask( + active_mask=(True, True, True, True, True, False, False, False) + ) + assert same_count_different_mask.active_token_count == _mask().active_token_count + assert same_count_different_mask.active_mask_sha256 != _mask().active_mask_sha256 + + +def test_padding_only_shard_is_constructible_for_the_identity_partial(): + sharding = _sharding( + vocab_shard_bounds=((0, QWEN3_REAL_VOCAB), (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB)), + ) + + assert sharding.local_vocab_start == 0 + assert sharding.vocab_shard_bounds[1] == (QWEN3_REAL_VOCAB, QWEN3_PADDED_VOCAB) + assert sharding.owner_rank(QWEN3_REAL_VOCAB - 1) == 0 From e6dbeefaba51f23176ddd72e70f257b81f070e1a Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Mon, 3 Aug 2026 06:09:14 +0800 Subject: [PATCH 5/6] docs(ws2): drop standalone design doc per review 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. --- docs/design/runtime-dispatch.md | 3 +- docs/design/ws2-tp-logprob-contract.md | 248 ------------------------- docs/operators/batch-invariant-logp.md | 4 +- rl_engine/kernels/logprob_contract.py | 34 +++- 4 files changed, 35 insertions(+), 254 deletions(-) delete mode 100644 docs/design/ws2-tp-logprob-contract.md diff --git a/docs/design/runtime-dispatch.md b/docs/design/runtime-dispatch.md index 84f29ec3..23c1586a 100644 --- a/docs/design/runtime-dispatch.md +++ b/docs/design/runtime-dispatch.md @@ -16,7 +16,8 @@ addition to platform priority, this path requires a backend capability descripto the requested role, dtype, TP/CP layout, padded-vs-real vocab masking, inactive-token support, vocab-domain LSE export, and deterministic TP merge semantics. Incompatible candidates produce explicit rejection reasons and are never used as an undeclared fallback. -See [WS2 TP-aware logprob contract](ws2-tp-logprob-contract.md). +The contract objects and their normative reduction semantics are documented in +`rl_engine.kernels.logprob_contract`. ## LogP Priority diff --git a/docs/design/ws2-tp-logprob-contract.md b/docs/design/ws2-tp-logprob-contract.md deleted file mode 100644 index cf9179b8..00000000 --- a/docs/design/ws2-tp-logprob-contract.md +++ /dev/null @@ -1,248 +0,0 @@ -# WS2 TP-Aware Logprob Contract - -Status: PR1 contract and dispatch metadata - -Tracking and shared contracts: - -- [#241: TP-aware deterministic logprob](https://github.com/RL-Align/RL-Kernel/issues/241) -- [#83: WS2 roadmap](https://github.com/RL-Align/RL-Kernel/issues/83) -- [#108: WS1 numerical contract](https://github.com/RL-Align/RL-Kernel/issues/108) -- [#111: WS2 cross-config alignment](https://github.com/RL-Align/RL-Kernel/issues/111) -- [#116: WS2 tolerance and drift-report format](https://github.com/RL-Align/RL-Kernel/issues/116) -- [Cross-config logprob drift contract](ws2_cross_config_logprob_drift_contract.md) - -## Scope - -This contract describes the logical inputs and deterministic reduction semantics for -selected-token log-probability under vocab-parallel tensor parallelism (TP): - -```text -selected_logp[t] = logits[t, target[t]] - logsumexp_vocab(logits[t, :]) -``` - -Under vocab-parallel TP each rank holds one vocabulary shard, so the vocabulary-wide -`logsumexp` requires a cross-rank reduction. This contract lets runtime dispatch reject a -backend whose numerical semantics do not match the requested layout. - -This PR1 layer does not shard tensors, launch a collective, merge `(max, sumexp)` partial -states, or implement a kernel. The single-GPU harness registration, the deterministic -vocab-parallel TP reference, and the cross-config integration belong to later PRs in #241. - -Context parallelism (CP) is a declared non-merge axis. CP partitions tokens, never the -vocabulary, so the logprob reduction spans TP vocab shards only. CP rank metadata is carried -for provenance and must never widen the merge. - -## Contract Objects - -`rl_engine.kernels.logprob_contract` defines: - -- `LogprobContract`: role, logits dtype, mask, sharding, reduction, output surface, and - LSE export, plus a rank-independent `cross_rank_fingerprint()`; -- `ShardingSpec`: per-rank vocab-shard bounds, padded-vs-real vocabulary, TP/CP rank - metadata, and target-token ownership; -- `MaskSpec`: active-token mask and ignore index; -- `ReductionSpec`: fixed `(max, sumexp)` merge semantics and the requested determinism - scope; -- `LogprobOutputSpec`: the output surface — fp32 selected logprob and fp32 vocab-domain - LSE, replicated across the TP group; -- `LogprobBackendCapability`: the layouts and semantics a backend explicitly supports, - including its mask modes and determinism scopes. - -Construction performs validation immediately. A structurally valid contract means that the -request is complete and internally consistent; it does not mean that an installed backend can -materialize it. - -`ShardingSpec.vocab_shard_bounds` lists every TP rank's half-open `[start, end)` vocab range -indexed by TP rank. The full table is required on every rank: it defines target ownership -and the fixed merge order without any collective, and it makes an incomplete or overlapping -partition a loud construction-time error instead of a silent runtime divergence. -`ShardingSpec.owner_rank(token_id)` resolves the unique owning rank for a real-vocab token -and rejects everything else. - -`padded_vocab_size` is the shard-covered (weight) vocabulary; `real_vocab_size` is the -tokenizer vocabulary. Padding columns occupy `[real_vocab_size, padded_vocab_size)` and must -be excluded from the logsumexp by any conforming implementation. The two sizes are equal -when the vocabulary is unpadded. - -Inactive tokens (prompt, padding, masked-out response positions) are excluded from every -drift aggregate and are exempt from the exactly-one-owner target gather; their targets may -legally hold `ignore_index`. `ignore_index` must not collide with the real vocabulary. - -## Qwen3-8B TP=2 BF16 Example - -```python -from rl_engine.kernels.logprob_contract import ( - LogprobContract, - MaskSpec, - ReductionSpec, - ShardingSpec, -) - -sharding = ShardingSpec( - tp_rank=0, - tp_world_size=2, - vocab_shard_bounds=((0, 76032), (76032, 152064)), - real_vocab_size=151936, - padded_vocab_size=152064, - cp_rank=0, - cp_world_size=2, -) - -contract = LogprobContract( - role="train", - dtype="bf16", - mask=MaskSpec( - num_tokens=8, - active_mask=(False, False, True, True, True, True, True, False), - ignore_index=-100, - ), - sharding=sharding, - reduction=ReductionSpec(), -) -``` - -Each rank owns one contiguous vocab shard; the 128 padding columns at the end of rank 1's -shard are outside the real vocabulary and never contribute to the logsumexp. The two leading -prompt tokens and the trailing padding token are inactive. - -## Reduction Semantics - -The only PR1 reduction contract is: - -```text -partial state: (local_max, local_sumexp), fp32 -merge: max_sumexp -merge_axis: tp_vocab -order: global_vocab_shard_index -transport: all_gather -downcast_at: final_write -engine: in_op_reference -``` - -Every rank first masks every local column whose global id lies in -`[real_vocab_size, padded_vocab_size)` to `-inf`, so padding never contributes to the -logsumexp, then computes `m_l = max(local_logits)` and `s_l = sum(exp(local_logits - m_l))` -in fp32. The partials travel by all-gather (collectives are transport only, never a -numerical reduction), and every rank merges in fixed global vocab-shard index order: - -```text -M = max_l(m_l) -S = sum_l(s_l * exp(m_l - M)) -LSE = M + log(S) -selected_logp = target_logit - LSE -``` - -The selected target logit comes from a masked single-owner gather: exactly one rank holds -each active token's target column. Downcast happens only at the final write. Because the -merge order is fixed by shard index, the result is deterministic and reproducible at every -TP degree by construction. Cross-degree bitwise equality (TP=2 equal to TP=1, the #241 PR 3 -acceptance target) requires one further condition: the entire reduction must follow a -global tile-level structure that is independent of TP partitioning — a fixed tile -decomposition of the vocabulary plus a fixed merge order and rescaling tree over those -tiles, identical at every TP degree, so that the TP degree only selects which rank computes -which tiles and never changes the floating-point grouping. A TP-degree-independent -decomposition inside each shard is not sufficient on its own, because shard boundaries -would still group the combines differently across degrees. Providing that global structure -is an obligation of the deterministic reference implementation; a backend without it is -still deterministic per degree, and its cross-degree drift is judged against the #108 -tolerance table instead. The contract expresses this distinction as -`ReductionSpec.determinism_scope`: `cross_tp_bitwise` (the #241 target and the default) -versus `fixed_topology`. A backend declares the scopes it honors in -`LogprobBackendCapability.determinism_scopes`, and dispatch rejects a backend that cannot -honor the requested scope — prose obligations are not enough; the guarantee is part of the -typed contract. - -A shard may lie entirely inside the padded region, and a row's local columns may all be -`-inf` after masking. The identity partial for these cases is defined as -`(m_l, s_l) = (-inf, 0)`: a partial with `s_l = 0` contributes nothing to the merge -regardless of its `m_l`, and implementations must use this identity directly rather than -evaluating `exp(-inf - (-inf))`, which would poison the merge with NaN. - -Averaging per-rank logsumexp values or letting a collective reduce numerically is never -conformant, at either determinism scope. - -The acceptable LSE and selected-token drift thresholds remain owned by #108, and drift -reports follow the #116 format. This contract does not introduce another tolerance table. -The selected-token metric remains the cross-config convention: - -```text -dlogp = training-side recomputed logp - rollout-side old logp -``` - -computed over active response tokens only. - -## Contract-Aware Dispatch - -Legacy callers continue to use `KernelRegistry.get_op()`. WS2 callers use: - -```python -result = kernel_registry.get_logprob_op(contract) -op = result.op -provenance = result.provenance -``` - -Dispatch considers only backends with a `LogprobBackendCapability`. It checks role, dtype, -TP/CP degree, padded-vs-real vocab masking, explicit active-mask support, vocab-domain LSE -export, and the requested determinism scope. An undeclared or incompatible backend is -skipped with an explicit rejection reason; there is no silent fallback. - -`requested_backend` accepts a case-insensitive policy keyword (`auto` | `production` | -`reference`; default `auto`) or an exact, case-sensitive stable backend id. The keywords -select an implementation tier; determinism is not a tier — it is requested through -`ReductionSpec.determinism_scope`, so `requested_backend="deterministic"` raises a loud -error instead of silently matching nothing. Strictness comes from the contract's -capability checks, not from the policy string. A backend id may never shadow a reserved -keyword; capability construction rejects that. The provenance `fallback` flag reports only -capability or load rejections of policy-eligible candidates — a candidate excluded by the -caller's own policy never counts as a fallback, even if it would also have failed -capability checks. - -WS2 dispatch resolves from its own candidate list, seeded from but decoupled from the legacy -`batch_invariant_logp` priority list: registering a TP-vocab backend for WS2 dispatch does -not change what legacy `get_op("batch_invariant_logp")` returns to WS1 callers. - -The current WS1 batch-invariant logp implementations are single-shard (TP=1) references: -they accept full-vocabulary logits with ignore-index masking but carry no vocab-shard -metadata, no padded-vs-real vocab distinction, and no public vocab-domain LSE export. Strict -WS2 requests therefore fail clearly today. The later deterministic vocab-parallel reference -becomes selectable through `KernelRegistry.register_logprob_backend(backend, capability)` -by declaring a capability that truthfully describes those features; no controller branch or -silent fallback is required. - -Successful dispatch provenance records: - -- requested and actual backend ids; -- platform and fallback status; -- prior candidate rejection reasons; -- the complete dispatch-relevant contract, including shard bounds, padded and real vocab - sizes, merge and output semantics, the explicit `cp_is_merge_axis: false` declaration, - and the active-mask digest (`active_mask_sha256`) — the mask's identity without its - per-token payload; -- the selected backend capability descriptor. - -### Distributed dispatch safety - -`get_logprob_op` resolves locally on each rank, so `requested_backend="auto"` is not -distributed-safe on its own: a load failure on one rank can resolve a different backend -than its peers, which for a collective-bearing implementation means divergent numerical -schedules or a deadlock. For `tp_world_size > 1` a caller must either request an exact -backend id or run a preflight agreement before any collective: all-gather the resolved -backend id together with `LogprobContract.cross_rank_fingerprint()` — a rank-independent -hash covering the shard-bounds table, vocab sizes, reduction/output semantics, and the -active-mask digest, excluding rank-local fields — and abort on any mismatch. Implementing -this preflight is an obligation of the #241 PR 3/PR 4 work; the backend invocation -protocol (how the contract and mask reach the implementation) is likewise defined there, -against this contract. - -## Validation - -Contract and dispatch behavior are covered by: - -```bash -python -m pytest tests/test_logprob_contract.py -q -``` - -The tests include Qwen3-8B TP=2 BF16 construction with padded vocab, the TP=1/2/4 sweep -shapes, incomplete/overlapping shard-bound rejection, owner-rank resolution, active-mask and -ignore-index validation, fp32-accumulation and merge-semantics enforcement, undeclared -backend rejection, no incompatible fallback, and JSON-compatible provenance. diff --git a/docs/operators/batch-invariant-logp.md b/docs/operators/batch-invariant-logp.md index b0671c40..4bca1f8c 100644 --- a/docs/operators/batch-invariant-logp.md +++ b/docs/operators/batch-invariant-logp.md @@ -64,8 +64,8 @@ remains unchanged. The backends above are single-shard (TP=1) references and do not yet export vocab-domain LSE or carry vocab-shard metadata, so they are declared incompatible with strict WS2 -requests instead of being selected as a silent fallback. See -[WS2 TP-aware logprob contract](../design/ws2-tp-logprob-contract.md). +requests instead of being selected as a silent fallback. The contract objects are +documented in `rl_engine.kernels.logprob_contract`. ## Benchmarks diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 6941c78e..9f6d72ba 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -100,8 +100,12 @@ class DeterminismScope(str, Enum): ``cross_tp_bitwise``: additionally bitwise-equal across TP degrees. This requires the entire reduction to follow a global tile-level structure that - is independent of TP partitioning (see the design doc); fixed shard-order - merging alone is not sufficient. + is independent of TP partitioning: a fixed tile decomposition of the + vocabulary plus a fixed merge order and rescaling tree over those tiles, + identical at every TP degree, so the TP degree only selects which rank + computes which tiles and never changes the floating-point grouping. + Fixed shard-order merging alone is not sufficient, because shard + boundaries would still group the combines differently across degrees. """ FIXED_TOPOLOGY = "fixed_topology" @@ -309,7 +313,31 @@ def active_mask_sha256(self) -> str: @dataclass(frozen=True) class ReductionSpec: - """Deterministic TP-vocab ``(max, sumexp)`` merge semantics.""" + """Deterministic TP-vocab ``(max, sumexp)`` merge semantics. + + Every rank first masks local columns whose global id lies in + ``[real_vocab_size, padded_vocab_size)`` to ``-inf`` (padding never + contributes to the logsumexp), then computes ``m_l = max(local_logits)`` + and ``s_l = sum(exp(local_logits - m_l))`` in fp32. Partials travel by + all-gather -- collectives are transport only, never a numerical + reduction -- and every rank merges in fixed global vocab-shard index + order:: + + M = max_l(m_l) + S = sum_l(s_l * exp(m_l - M)) + LSE = M + log(S) + selected_logp = target_logit - LSE + + The selected target logit comes from a masked single-owner gather; + downcast happens only at the final write. The identity partial for a + padding-only shard, or a row whose local columns are all ``-inf`` after + masking, is ``(m_l, s_l) = (-inf, 0)``: a partial with ``s_l = 0`` + contributes nothing to the merge regardless of its ``m_l``, and + implementations must use this identity directly rather than evaluate + ``exp(-inf - (-inf))``, which would poison the merge with NaN. Averaging + per-rank logsumexp values, or letting a collective reduce numerically, is + never conformant at either determinism scope. + """ merge: LogprobMerge = LogprobMerge.MAX_SUMEXP merge_axis: MergeAxis = MergeAxis.TP_VOCAB From 878ba88a44e5f9ac24d1a38b4536a60c02721602 Mon Sep 17 00:00:00 2001 From: ryankert01 Date: Mon, 3 Aug 2026 06:21:30 +0800 Subject: [PATCH 6/6] style(ws2): align comment density with sibling kernel modules 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. --- rl_engine/kernels/logprob_contract.py | 65 +++++++++------------------ rl_engine/kernels/registry.py | 19 +++----- 2 files changed, 26 insertions(+), 58 deletions(-) diff --git a/rl_engine/kernels/logprob_contract.py b/rl_engine/kernels/logprob_contract.py index 9f6d72ba..b0bcd7ab 100644 --- a/rl_engine/kernels/logprob_contract.py +++ b/rl_engine/kernels/logprob_contract.py @@ -29,16 +29,10 @@ _EnumT = TypeVar("_EnumT", bound=Enum) -# Dispatch policy keywords accepted by KernelRegistry.get_logprob_op; a stable -# backend id must never shadow one of these, or it becomes unselectable by id. -# "deterministic" stays reserved even though it is no longer a policy: -# determinism is expressed through DeterminismScope, and requesting it as a -# policy is a loud error rather than a silent id mismatch. +# Policy keywords accepted by KernelRegistry.get_logprob_op; a backend id must +# never shadow one of these, or it becomes unselectable by id. RESERVED_DISPATCH_POLICIES = frozenset({"auto", "production", "reference", "deterministic"}) -# Implementation tiers a backend can declare. Determinism is deliberately a -# separate axis (DeterminismScope): a backend can be a deterministic reference, -# a deterministic production implementation, or a non-deterministic production -# implementation. +# Backend tiers; determinism is a separate axis (DeterminismScope). IMPLEMENTATION_KINDS = frozenset({"production", "reference"}) @@ -58,12 +52,7 @@ class LogprobDType(str, Enum): class LogprobMerge(str, Enum): - """Merge primitive for per-shard partial states. - - Every rank contributes ``(local_max, local_sumexp)`` computed in the - accumulation dtype; the merged result is - ``M = max(m_l)``, ``S = sum(s_l * exp(m_l - M))``, ``LSE = M + log(S)``. - """ + """Merge primitive for per-shard ``(local_max, local_sumexp)`` partials.""" MAX_SUMEXP = "max_sumexp" @@ -115,11 +104,9 @@ class DeterminismScope(str, Enum): class MaskMode(str, Enum): """How a backend consumes inactive-token information. - ``explicit_active_mask``: the backend honors an arbitrary active-token - mask. ``ignore_index``: the backend only recognizes inactive tokens whose - target id equals ``ignore_index``. The contract permits inactive targets - that do NOT hold ``ignore_index``, so an ignore-index-only backend cannot - serve a contract with inactive tokens. + The contract permits inactive targets that do not hold ``ignore_index``, + so an ``ignore_index``-only backend cannot serve a contract with inactive + tokens. """ EXPLICIT_ACTIVE_MASK = "explicit_active_mask" @@ -161,15 +148,11 @@ class ShardingSpec: """Logical vocab-parallel TP ownership for one logprob invocation. ``vocab_shard_bounds`` lists every TP rank's half-open ``[start, end)`` - vocab range indexed by TP rank. The full table is required on every rank: - it defines target-token ownership and the fixed global-shard-index merge - order without any collective, and makes an incomplete partition a loud - construction-time error instead of a silent runtime divergence. - - ``padded_vocab_size`` is the shard-covered (weight) vocabulary; - ``real_vocab_size`` is the tokenizer vocabulary. Padding columns occupy - ``[real_vocab_size, padded_vocab_size)`` and must be excluded from the - logsumexp by any conforming implementation. + vocab range, indexed by rank; the full table is required on every rank and + must form a contiguous ``[0, padded_vocab_size)`` partition. + ``padded_vocab_size`` is the shard-covered (weight) vocabulary, + ``real_vocab_size`` the tokenizer vocabulary; padding columns occupy + ``[real_vocab_size, padded_vocab_size)``. """ tp_rank: int @@ -267,11 +250,10 @@ def owner_rank(self, token_id: int) -> int: @dataclass(frozen=True) class MaskSpec: - """Active-token ownership for one logprob invocation. + """Active-token mask and ignore index for one logprob invocation. - Inactive tokens are excluded from every drift aggregate and are exempt - from the exactly-one-owner target gather; their targets may legally hold - ``ignore_index``. + Inactive tokens are excluded from drift aggregates and from the + single-owner target gather; their targets may legally hold ``ignore_index``. """ num_tokens: int @@ -377,12 +359,8 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class LogprobOutputSpec: - """Output surface every conforming backend must produce. - - Selected logprob and vocab-domain LSE are fp32 and replicated across the - TP group; the ``downcast_at: final_write`` rule applies to any consumer - downcast after these outputs, never inside the reduction. - """ + """Output surface every conforming backend must produce: fp32 selected + logprob and fp32 vocab-domain LSE, replicated across the TP group.""" selected_logp_dtype: LogprobDType = LogprobDType.FP32 lse_dtype: LogprobDType = LogprobDType.FP32 @@ -466,10 +444,8 @@ def to_dict(self) -> dict[str, Any]: "determinism_scope": self.reduction.determinism_scope.value, "cp_is_merge_axis": False, } - # The per-token mask is deliberately summarized: provenance exists for - # logging/serialization and the raw mask would dominate its size. The - # digest keeps the mask *identity* observable, so two masks with the - # same active count still produce distinguishable provenance. + # The digest stands in for the raw per-token mask, which would + # dominate the provenance size. mask = { "num_tokens": self.mask.num_tokens, "active_token_count": self.mask.active_token_count, @@ -614,8 +590,7 @@ def incompatibilities(self, contract: LogprobContract) -> tuple[str, ...]: contract.mask.active_token_count != contract.mask.num_tokens and MaskMode.EXPLICIT_ACTIVE_MASK not in self.mask_modes ): - # The contract does not require inactive targets to hold - # ignore_index, so ignore-index-only masking is insufficient. + # Inactive targets need not hold ignore_index (see MaskMode). reasons.append("explicit active-token masking is unsupported") if contract.export_lse and not self.exports_vocab_lse: reasons.append("vocab-domain LSE export is unsupported") diff --git a/rl_engine/kernels/registry.py b/rl_engine/kernels/registry.py index b17acf3e..7093b320 100644 --- a/rl_engine/kernels/registry.py +++ b/rl_engine/kernels/registry.py @@ -175,13 +175,9 @@ def __init__(self): self._instance_cache: Dict[str, Any] = {} self._failed_backends: Set[str] = set() - # These descriptors report what the existing WS1 batch-invariant logp - # implementations actually support: single-shard (TP=1) logits with - # ignore-index masking, no vocab-shard metadata, no padded-vs-real - # vocab distinction, and no public vocab-domain LSE export. A strict - # WS2 request is rejected with explicit reasons until the deterministic - # vocab-parallel TP reference backend lands (issue #241 PR 3) instead - # of silently selecting an incompatible fallback. + # Truthful descriptors for the existing WS1 batch-invariant logp + # implementations: single-shard (TP=1), ignore-index masking only, no + # vocab-shard metadata, no vocab-domain LSE export. common_logprob_roles = frozenset({LogprobRole.TRAIN, LogprobRole.INFER}) common_logprob_dtypes = frozenset({LogprobDType.BF16, LogprobDType.FP16, LogprobDType.FP32}) base_logprob_capabilities = { @@ -336,12 +332,9 @@ def __init__(self): self._adjust_priority_for_hardware() self._adjust_priority_from_env() - # WS2 contract-aware dispatch owns its candidate list. It is seeded - # from the legacy batch_invariant_logp priority (after hardware/env - # adjustments) but deliberately decoupled afterwards: registering a - # TP-vocab backend for WS2 dispatch must not change what legacy - # get_op("batch_invariant_logp") returns to WS1 callers, and vice - # versa. + # WS2 dispatch owns its candidate list, seeded from the legacy + # batch_invariant_logp priority but decoupled afterwards: neither + # path's registrations may affect the other. self._logprob_candidates: Dict[str, list] = { platform: list(ops.get("batch_invariant_logp", [])) for platform, ops in self._priority_map.items()