[WS2][Cross-config][Attention] Integrate the Qwen3-8B TP=2 CP=2 BF16 Attention setting to cross-config alignment framework - #263
Open
zhangj1an wants to merge 15 commits into
Conversation
Signed-off-by: inaniloquentee <3051000145@qq.com>
Signed-off-by: inaniloquentee <3051000145@qq.com>
Signed-off-by: inaniloquentee <3051000145@qq.com>
# Conflicts: # rl_engine/kernels/registry.py
# Conflicts: # docs/operators/attention.md
…ign#235 PR4) Wire the CP attention path into the cross-configuration planner/runtime for the Qwen3-8B TP=2 CP=2 BF16 target. The PR4 criterion "rollout and training descriptors bind to the same semantic attention contract" cannot hold literally: training runs full-sequence prefill over a CP-sharded sequence while rollout runs vLLM paged-KV chunked prefill, so the two AttentionContract instances always differ. Binding is therefore split into three tiers -- identity must match bit for bit, reduction semantics must match each other and the WS2 mandate, and materialization differences are recorded and measured rather than rejected. reduction.engine stays in the recorded tier so a Transformer Engine merge oracle on one side does not fail the binding; reduction.order and acc_dtype stay in the semantic tier because that is the WS2 claim. Also adds the first two framework-shaped RuntimeMaterializer implementations. Before this the only one was CpuSmokeMaterializer over a synthetic CPU model, and every named scenario was planning-only. Neither adapter imports megatron or vllm, so the binding rules run on CPU in CI. Determinism is probed on both sides and compared, because the two frameworks mean different things by it: Megatron asserts NCCL_ALGO and leaves TF32 and BF16 reduced-precision reduction unmanaged, while vLLM hard-sets ten NCCL variables and disables both. Mismatches in NCCL_ALGO, NCCL_PROTO and CUBLAS_WORKSPACE_CONFIG are blocking; the rest are recorded. Fixes a latent break on the way: the planner normalizes dtype knobs to torch spellings (bfloat16) while AttentionDType uses short ones (bf16), so passing a normalized knob into the enum raised. Stacked on RL-Align#236 (attention contract) and RL-Align#238 (deterministic CP reference), on top of RL-Align#230 (cross-configuration framework). Part of RL-Align#235 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw
PR4) Three fields could differ between the two sides without the binding noticing. dtype was in no tier at all, so a BF16 rollout could bind to an FP16 training pass and produce a drift number attributable to nothing. It joins the semantic tier, with allow_dtype_difference for the RL-Align#235 PR5 sweep that deliberately scores BF16 against an FP32 reference. batch_size was likewise unchecked. Batch invariance is a claim about results not changing with batch makeup, so two sides scoring different batches are not comparable and it belongs to identity. split_kv_policy has no field in the RL-Align#236 contract, so it only reached side_configs and never took part in binding. Callers now pass it through rollout_recorded_extra / training_recorded_extra so the difference is at least visible in provenance; it can move into the contract once RL-Align#236 grows the field. Part of RL-Align#235 Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Q3Ar3z9fHEBFQQHddSEMaw
zhangj1an
requested review from
Flink-ddd,
KJLdefeated,
bitborne and
inaniloquentee
as code owners
August 4, 2026 14:56
|
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 |
Open
8 tasks
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
2 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.
Latest Status [Aug 4 2026]
I will update the table below after I fixed upstream PR for cases marked as ❌ below.
Also, I will update the config for RoPE in this PR. (it would not cause training / inference inconsistency)
Motivation
This PR implements #235 , PR4.
This branch is simultaneously based on #230 and #235. #236 and #238 should be merged before this PR is merged.
Dependency Stack
This PR only adds 4 knobs that will affect attention kernel results, according to Chutian's comments.
attention.reduction_orderReductionSpec.orderReductionOrder.ARRIVAL; #238 add an arrival-order reference variant; #235 PR5 run the sweepattention.split_kv_policysplit_kv_policytoAttentionContract+ supported-k list toAttentionBackendCapability; #235 PR7 write and align the k = 2/4/8 variantsbatch.sizeAttentionContract.batch_sizetraining.compute_dtype/rollout.dtypeAttentionContract.dtypeallow_dtype_differenceWe do not turn on CP because attention do not use TP.