Consolidate sampling-transform scoring on one shared chunked core; implement top-k/top-p/min-p replay for generic ce_modes

[qywu]

Addresses the scoring-path half of #71 (points 1, 2, 3, 5 of the proposed direction) and closes the #60 gap: a generic (non-exact-family) model can now train on rollouts sampled with top_k/top_p/min_p and per-row temperature, getting logprobs on the filtered support, under ordinary body TP.

Refined plan (updating the issue's direction)

The issue's five points remain the plan; this PR deliberately takes the scoring path and leaves point 4 (capability negotiation across the 78 _exact_contract_family branch sites — PP byte contract, weight sync, checkpointing) to a follow-up PR series. Two inputs for that follow-up surfaced while inventorying the branch surface:

• _qwen3_dense_exact_contract never becomes an _exact_contract_family value, so under PP an exact dense Qwen3 model does not get _pp_exact_boundary_contract (only DSV4 is explicitly re-added in engage_pp_byte_contract), even though it is pinned to bi_fused and BI trunk linears.
• The loss layer already keys off module marks (_glm52_exact_tp16_lm_head, _dsv4_exact_tp8_lm_head), not the family — this PR leans into that: once the generic path is capability-complete for transform replay, the loss layer no longer needs lane membership at all.

What this PR does

1. One documented sampling-transform program (issue point 1). xorl/ops/exact_sampling_transforms.py is now declared the XoRL-wide replay contract for every lane, not an exact-lane-only artifact. The docstring pins: temperature-first at the lane's dtype boundary, stable token-ID tie-break, inclusive top-p crossing, min-p against the original row maximum, and top_p == 1.0 as the exact identity despite fp32 cumulative overshoot (the question raised on #60; already tested at tests/ops/test_exact_sampling_transforms.py::test_identity_top_p_one_is_full_support_despite_fp32_cumulative_overshoot). It also states explicitly that stock SGLang/FlashInfer filters do not guarantee these rules, so exact replay requires a sampler that implements this program.

2. One shared chunked scoring core (issue point 2). New xorl/ops/loss/sampling_transform_ce.py: chunked projection → temperature → pinned support program → selected logprob, with a recompute closed-form VJP, saving only per-row scalars (never a [tokens, vocab] activation). The genuinely lane-specific piece — how chunk logits are computed — is an injected policy:

• generic modes inject a plain GEMM (fp32 per lm_head_fp32);
• bi_fused injects the batch-invariant kernels (head_v2_full_logits_with_lse / bi_lm_head_full_logits + exact_temperature_scale_fp32_logits), so its forward values keep matching serving bytes; its duplicated chunk loops and filtered backward are deleted (bi_fused_lm_head.py −362/+94 net).
• The GLM-5.2 and DSV4 exact heads keep their QLoRA-surrogate VJPs but now consume the shared temperature validation/normalization (was 5 copies) and the shared reference-VJP trio (was 2 copies). Their dead legacy-call autograd shims are removed.

Ordinary vocabulary-parallel body TP is handled inside the core: chunk-local logits are gathered in rank order, hidden gradients are summed across vocab ranks, weight-shard gradients stay local.

3. Transforms implemented on the generic modes (issue point 3). per_token_ce.py:274/:337 (NotImplementedError: top-k/top-p/min-p replay is supported only by exact LM-head modes) and the three per-row-temperature rejections are gone. All of eager / compiled / quack_linear / fused_quack now score transform-carrying requests through the shared core, with and without body TP.

4. One loud rejection in one place (issue point 5). The single remaining unsupported combination — transforms through an FP8 lm_head module — raises one actionable message (SAMPLING_TRANSFORM_LM_HEAD_MODULE_ERROR) naming the fix (lm_head_fp32: true). A replayed token outside its row's current support scores -inf logprob (+inf CE) with zero gradient — no path silently scores against unfiltered support.

Acceptance criteria status

• Generic model trains on top-k/top-p/min-p rollouts with filtered-support logprobs under ordinary body TP (tests/distributed/test_vocab_parallel_ce.py::check_sampling_transform_replay — CE and both grads bit-match the dense reference at TP2)
• One documented sampling-transform program with divergences written down and tested
• The chunked scoring core exists once; bi_fused consumes it directly, GLM-5.2/DSV4 consume the shared program + helpers (their kernel-exact QLoRA VJPs remain theirs by design)
• No silent unfiltered-support scoring (out-of-support ⇒ +inf CE, zero grad, tested)
• Branch-site reduction on _exact_contract_family — follow-up (point 4)

Not in this PR (follow-ups)

• Capability negotiation replacing family gating at the remaining branch sites (issue point 4), including the _qwen3_dense_exact_contract PP hole above.
• Physical-PP transform plumbing (model_runner._pp_sampling_transform_kwargs keeps its loud error; the per-token metadata queue that temperature uses needs the same treatment for the filter fields).
• The bi_fused dedicated-LM-head-TP function (_BiFusedVocabParallelPerTokenCE) keeps its own row-owner broadcast structure; it now shares the gather + scoring helpers but not the core loop.

Self-review pass (second commit)

An adversarial review of the diff surfaced and the second commit fixes: a stale 7-arg autograd call in tests/distributed/test_glm52_exact_lm_head_qlora_collectives.py missed by the legacy-shim removal; a validation gap where over-length transform metadata was silently truncated (now validate_sampling_transform_rows, shared by the core and both bi_fused entries, with a regression test); a backward perf regression for the common temperature-only bi_fused configuration (restored as a single-pass vocabulary-chunked backward in the core — write-once weight-gradient slices, no full-row logits, no TP gather in backward); addmm(out=) accumulation in the filtered backward instead of a per-chunk [V, H] fp32 temporary; the missing <= 0 clamp in _plain_selected_score; and the scalar (pre-GEMM) vs per-row-tensor (post-GEMM) temperature convention split is now documented.

Known perf headroom deliberately left for later: the filtered TP path gathers full-vocab logits per row chunk in forward and backward and re-sorts on every rank (this path previously did not exist at all — correctness first), and gather_vocab_shards could move to xorl.utils.dist_utils alongside unifying the glm5/dsv4 private rank-order gathers.

Third commit: exact-head deduplication (formerly #75, folded in)

The GLM-5.2 and DSV4 exact heads ran the same program skeleton in near-duplicate. New xorl/models/transformers/exact_lm_head_shared.py owns the skeleton — one ExactLmHeadFunction autograd boundary (row ownership injected as an ExactHeadRowPlan: replicated / equal TP16 blocks / ragged TP8 blocks), one filtered + one unfiltered surrogate-VJP plumbing loop, the serving LoRA A/B SGEMM choreography, the rank-order collectives, and one TP-group geometry checker. Each family keeps only its contract-bearing pieces: base projection kernel (head_v2 FP32 vs F.linear BF16), temperature dtype boundary, native selected-score kernel, surrogate VJP formulation, and geometry validation. Forward value paths are byte-preserved — the same kernels run in the same order with the same dtype boundaries. gather_vocab_shards moved to xorl.utils.dist_utils. Net −575 lines (glm5 head 1298→847, dsv4 head 836→612, +369 shared).

Fourth commit: review fixes on the dedup

An adversarial review of the dedup commit found (and the fourth commit fixes) a real bug it introduced: the shared autograd boundary inherited GLM's FP32 selected-logprob assertion, while DSV4's pinned value program is BF16 end-to-end — the DSV4 exact head raised on its first real forward, masked in CI by FP32-returning fake components. The expected dtype is now part of the family contract (component.logprob_dtype), the DSV4 fakes pin the real BF16 contract, and a mismatch test covers the gate. Also fixed: the shared backward revalidates the TP group before its first collective (deterministic contract error instead of an NCCL hang on a rebound/destroyed group), require_equal_nonzero_row_count reuses rank_order_row_counts, and ExactHeadRowPlan documents its no-tensor-capture contract. Noted, not changed: DSV4's collapsed legacy error strings ("DSV4 exact lm-head TP8 rank/order mismatch") are replaced by the checker's more specific per-condition messages; the pad-to-max row-gather choreography still has siblings in vocab_parallel_reverse_kl/opd_loss/model_runner (follow-up); the ~7 per-row gathers per scored token predate this PR and can now be batched in one place.

Testing

• New tests/ops/loss/test_sampling_transform_ce.py (12 tests, CPU): filtered-support parity with the dense program reference across all four generic modes, gradient parity with full autograd, out-of-support ⇒ +inf/zero-grad, per-row temperature vs scaled CE, lm_head_fp32 convention, identity metadata collapses to the untouched fast path (bit-equal), FP8-module rejection.
• tests/distributed/test_vocab_parallel_ce.py gains the body-TP transform-replay check (2×H100: CE/grad_h/grad_w max err 0.0 vs dense reference).
• Existing surfaces re-run green: tests/ops/test_exact_sampling_transforms.py (14), tests/ops/loss/ + exact-head + wiring tests (77 total), GPU tests/ops/test_bi_fused_lm_head.py (12), tests/distributed/test_bi_fused_lm_head_tp.py, test_fused_linear_logprob_tp.py, and the model-runner/trainer transform-relay tests (66).

References are against main @ e12b12b.

[broly-code-security-scanner]

Broly Security Scan

Note

✅ Clean scan
No vulnerabilities detected in this PR.

Note

Re-scan this PR anytime with /broly scan — useful after /broly undismiss, or to refresh findings without a new push.

Broly — SAST (zai-org/GLM-5.2) · Secrets · SCA · IaC · GH Actions · Base Images · Supply Chain Threats · Exploit Chains · Adversarial Verification

We're continuously improving Broly's accuracy and finding quality — your feedback is valuable. False positives, missed findings, bugs, and feature requests all welcome.