feat(logprob): add deterministic ROCm vocab-parallel path - #328
Open
hihaluemen wants to merge 58 commits into
Open
feat(logprob): add deterministic ROCm vocab-parallel path#328hihaluemen wants to merge 58 commits into
hihaluemen wants to merge 58 commits into
Conversation
Implements PR 1 of issue RL-Align#241: a typed contract for vocab-parallel selected-token logprob, mirroring the WS2 attention contract pattern. - rl_engine/kernels/logprob_contract.py: LogprobContract, ShardingSpec (per-rank vocab shard bounds, padded-vs-real vocab, TP/CP rank metadata, owner_rank resolution), MaskSpec (active-token mask, ignore_index), ReductionSpec (fp32 (max, sumexp) merge in fixed global vocab-shard index order, all-gather transport, CP declared a non-merge axis), and LogprobBackendCapability. - KernelRegistry.get_logprob_op(contract): contract-aware dispatch that only selects backends with a declared capability; incompatible or undeclared candidates are rejected with explicit reasons and never used as a silent fallback. Existing WS1 batch-invariant logp backends are declared truthfully as single-shard references, so strict WS2 requests fail loudly until the deterministic vocab-parallel TP reference (PR 3) lands. Legacy get_op() behavior is unchanged. - Design doc, runtime-dispatch and operator doc updates, and CPU-safe contract/dispatch tests covering the Qwen3-8B TP=2 BF16 target and the TP=1/2/4 sweep shapes. Tolerance values remain owned by RL-Align#108.
- docs: correct the TP-invariance claim — fixed merge order gives determinism per TP degree; cross-degree bitwise equality additionally requires a TP-degree-independent local tile decomposition (PR 3 obligation), otherwise RL-Align#108 tolerances apply - contract: store backend_id stripped so id-based dispatch matches; summarize the active mask in to_dict() provenance instead of copying every per-token boolean; sort __all__ per RUF022 - registry: add public register_logprob_backend() seam for PR 3 and tests; delegate _platform() to _platform_for_device(None); reuse _get_or_create_backend() in get_op so WS2 and legacy dispatch share one cache/blacklist code path - tests: use the registration seam instead of poking private state, pin _even_bounds' last bound for non-divisible vocabularies, assert candidate-list decoupling in both directions, cover registration replace semantics and backend_id normalization
- docs: state that cross-TP bitwise equality needs a global tile-level merge structure independent of TP partitioning (per-shard tiles alone leave different grouping at shard boundaries), and that padded columns are masked to -inf before the local (max, sumexp) partials - registry: scope logprob capabilities per platform so the same backend enum can declare different support on cuda/rocm/cpu; validate the platform argument of register_logprob_backend against known platforms - contract: derive IMPLEMENTATION_KINDS from RESERVED_DISPATCH_POLICIES and use it for the kind check; wrap non-iterable roles/dtypes in LogprobContractError for consistent error handling - tests: cover per-platform capability scoping, unknown-platform rejection, and non-iterable roles/dtypes
…typed contract Address external review: the cross-TP bitwise guarantee lived only in prose, so a fixed-topology-deterministic backend could pass dispatch as fully conformant. - DeterminismScope (fixed_topology | cross_tp_bitwise): requested via ReductionSpec (default cross_tp_bitwise, the RL-Align#241 PR 3 target), declared per backend via determinism_scopes, enforced by dispatch; replaces the deterministic_tp_merge bool - MaskMode (explicit_active_mask | ignore_index) replaces supports_inactive_tokens: the contract permits inactive targets that do not hold ignore_index, so ignore-index-only backends are rejected for contracts with inactive tokens - LogprobOutputSpec pins the output surface: fp32 selected logprob and fp32 vocab LSE, replicated across the TP group - implementation_kind is now a tier (reference | production); determinism is no longer conflated with it, and requesting "deterministic" as a policy raises a loud error pointing at determinism_scope - fallback provenance: policy evaluation now precedes capability checks, so a candidate excluded by the caller's own policy never counts as a fallback even when it also lacks capabilities - docs: define the (-inf, 0) identity partial for padding-only or all--inf shards; document that requested_backend="auto" is not distributed-safe and specify the preflight fingerprint agreement - LogprobContract.cross_rank_fingerprint(): rank-independent identity for that preflight; provenance now records active_mask_sha256 so masks with equal active counts remain distinguishable
Fold the normative reduction semantics (padded-column masking, fp32 (max, sumexp) merge formulas, the (-inf, 0) identity partial, and the cross-TP tile-structure requirement) into the ReductionSpec and DeterminismScope docstrings, and repoint the runtime-dispatch and batch-invariant-logp doc references at the module. The contract summary moves to the PR description.
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.
The guard added per review rejects requested_backend="auto" whenever tp_world_size > 1, so TP-sharded dispatch tests now name an explicit policy and auto-policy tests use TP=1 contracts. Add coverage for the guard itself and document the restriction in get_logprob_op. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012PyjQEqDJwy9Cos4Sb9QBK
…logprob-rocm # Conflicts: # csrc/ops.cpp # setup.py
hihaluemen
requested review from
EthanZero2Hero,
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
August 21, 2026 20:58
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
hihaluemen
force-pushed
the
work/ws2-logprob-rocm
branch
2 times, most recently
from
August 21, 2026 21:08
1a2eb62 to
df1e637
Compare
hihaluemen
force-pushed
the
work/ws2-logprob-rocm
branch
from
August 21, 2026 21:11
df1e637 to
f79602c
Compare
20 tasks
zhangj1an
added a commit
that referenced
this pull request
Aug 24, 2026
…patch Adds the dispatch seam and Vime adapter that make the strict ROCm attention core reachable as an explicit, fail-closed backend, mirroring the WS2 logprob provider in #328. Registry: - register_attention_backend / get_attention_op over AttentionContract, with a candidate list kept separate from the legacy attn/attention priority maps and empty on every platform by default. A WS2 caller can therefore never be served by an SDPA-shaped wrapper that does not export attention-domain LSE, and the strict core can never be selected by a legacy get_op caller. - aiter.rocm.ck_dense_mha is registered only when aiter.ops.mha genuinely imports and exposes mha_fwd/mha_bwd. An explicit request for it fails loudly when the vendor stack is absent instead of degrading to a different backend. Vime adapter (rl_engine/integrations/vime/attention.py): - Structural request in, (out, lse) plus provenance out; RL-Kernel never imports Vime, and native fallback is signalled through an attention_provider_unavailable marker rather than an imported type. - Each logical batch row is materialized on its own. Raw AITER mha_fwd is batch-composition sensitive in BF16 at some shapes (S=256 B=4, S=512 B=2/B=4; up to 1.5625e-02) while invariant at others, so batching would otherwise change the bits for a subset of shapes only. - Position identity is derived and validated, so a training-side full sequence and a rollout-side chunk provably describe the same tokens. - CP>1 and decode fail closed rather than being served by a core that does not own the cross-rank merge or KV-cache identity. Contract: - AttentionContract.cross_rank_fingerprint() for cross-rank preflight, mirroring LogprobContract; auto dispatch stays rejected under CP>1. Benchmarks (MI300X, gfx942, torch 2.12.0+rocm7.14, BF16, Qwen3-8B heads): forward is within 4% of SDPA and ~1.9x faster than the Triton FlashAttention backend at S=4096. The deterministic backward costs ~2.4x time and scales O(S^2) in memory (16.8 GiB at S=4096); that cost is AITER's mha_bwd, not the integration - raw AITER peaks at 16657 MiB against the provider's 16722 MiB. Tests: 38 new cases; the full suite is unchanged against the pre-change baseline (27 failed / 224 skipped / 607 errors both before and after, all from the unbuilt native extension), with passed rising 1068 -> 1102. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QfN3b2ep5DvxVDir7J36j3
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add deterministic ROCm backend for WS2 vocab-parallel logprob path (#241)
This adds a ROCm HIP path for the WS2 vocab-parallel logprob implementation. The native part is deliberately small: it only handles the rank-local vocab-tile statistics. The existing TP merge order and logprob/entropy contract are left alone, which keeps the training and inference paths on the same definition.
Changes
deterministic_logp_tile_statsfor FP32 max/sumexp partials. It uses a fixed block reduction and noatomicAdd, and the same source is used for CUDA and HIP builds.rocm-vocab-parallel-logp-ws2backend. It is registered only when the ROCm extension and symbol load successfully. An explicit request for the native backend fails immediately if the extension is missing; it will not silently turn into the PyTorch reference path.Benchmark results
Operator-only benchmark — seeded logits, no checkpoint, tokenizer, or serving engine.
Qwen3 vocabulary
V=151936split into 64 tiles of 2374 columns. Tables show 2048 tokens, BF16.mi300xh100cpuSingle device (BF16, 2048 tokens)
Distributed (BF16, 2048 tokens, forward / forward+backward median ms)
Full report ·
Raw JSON