Consolidate the exact and generic scoring paths; implement top-k/top-p/min-p on the generic path

[qywu]

Problem

XoRL currently maintains two parallel scoring/serving programs:

• the exact path — the family-gated bit-exact train/serve value programs (GLM-5.2, Qwen3.5 dense, Qwen3.5 MoE, DSV4), selected by resolve_exact_contract_family in src/xorl/models/exact_contract.py
• the generic path — everything else: ce_mode in {eager, compiled, quack_linear, fused_quack} plus the ordinary HF-derived model families (llama3, qwen2, qwen3, qwen3_moe, glm4_moe, gpt_oss, minimax_m3, nemotron_h, olmo2, deepseek_v3)

The split is not a clean layering. It is a fork that has to be re-resolved by hand on every change, and it leaves the generic path missing capabilities the request schema already accepts.

1. The branch surface keeps conflicting

resolve_exact_contract_family returns None for generic models, and the resulting stamp (config._exact_contract_family) plus its sibling predicates are branched on at 78 sites across 13 files:

src/xorl/trainers/model_builder.py
src/xorl/models/auto.py
src/xorl/models/layers/rope.py
src/xorl/models/layers/moe/moe_block.py
src/xorl/models/transformers/{glm5/modeling_glm5.py,glm5/indexer.py,glm5/support.py}
src/xorl/models/transformers/{qwen3_5/modeling_qwen3_5.py,qwen3_5_moe/modeling_qwen3_5_moe.py}
src/xorl/distributed/pp_byte_contract.py
src/xorl/server/runner/checkpoint/manager.py
src/xorl/server/weight_sync/handler.py

Every cross-cutting change (RoPE, MoE reduce order, PP byte boundary, weight sync, checkpointing) has to be written twice and reconciled at each of those branch points. That is the mechanical source of the repeated conflict churn — e.g. #65 had to resolve 71 replay conflicts plus 13 more on the final rebase, almost all of it in code that exists in near-duplicate exact/generic forms.

The LM-head scoring program itself is triplicated:

file	lines
src/xorl/models/transformers/glm5/exact_lm_head_qlora.py	1396
src/xorl/models/transformers/deepseek_v4/exact_lm_head.py	923
src/xorl/ops/loss/bi_fused_lm_head.py	903

All three implement the same conceptual op (chunked projection → fixed-order LSE → exact_sampling_support → selected logprob + VJP). Only the support predicate in src/xorl/ops/exact_sampling_transforms.py is genuinely shared; the chunking, temperature boundary, TP reduction, and autograd wrapper are copies that can and do drift.

2. Model support is gated by lane, not by capability

Anything a model needs from the exact lane requires being in an exact family — there is no way to opt a generic model into an individual capability. Concretely, src/xorl/ops/loss/per_token_ce.py rejects per-row logprob_temperature outside exact LM-head modes at lines 276, 341, 372, 383, and rejects sampling transforms entirely (below). So the ~10 generic families cannot do behavior-policy replay at all, even though nothing about the math is family-specific.

3. Generic path does not implement top-k / top-p / min-p post-processing (#60)

This is the sharpest instance of (2), and it is a plumbing/implementation mismatch, not a missing feature request:

• The orchestrator already accepts the sampler's normalized transforms and materializes them as token-aligned replay metadata — NORMALIZED_SAMPLING_METADATA_FIELDS at src/xorl/server/orchestrator/packing.py:81-86, expanded at :196-227.

• Collation, padding, PP, and shared-prefix repack all carry the fields: src/xorl/data/collators/packing_concat_collator.py:97, src/xorl/distributed/sync_padding.py:26, src/xorl/ops/shared_prefix/repack.py:48, src/xorl/server/runner/model_runner.py:5843-5864.

• Then the trainer hard-fails: src/xorl/ops/loss/per_token_ce.py:274 (fused_quack) and :337 (eager/compiled/quack_linear) both raise

  NotImplementedError: top-k/top-p/min-p replay is supported only by exact LM-head modes
  

• Under physical PP the same request shape raises a different error from a different place: model_runner.py:7387-7394, "Physical PP received top-k/top-p/min-p metadata before the exact sampling-transform loss helpers were composed".

So the only route to top-k/top-p replay is ce_mode="bi_fused", which fails closed on constraints most generic runs violate:

• lm_head_fp32: true required (per_token_ce.py:304-307)
• CUDA + BF16 hidden and weight required (bi_fused_lm_head.py:779-781)
• no FP8 lm_head module (per_token_ce.py:302-303)
• TP only through the dedicated vocabulary-sharded LM-head group — ordinary body TP is refused (per_token_ce.py:44-62, :321-324)

Net effect: a rollout that used top_p < 1 or top_k < vocab against a generic model either crashes at loss time or, if the sampler metadata never gets relayed, silently trains on unfiltered-support logprobs — i.e. wrong importance ratios with no error.

There is also an unresolved semantics question on top: exact_sampling_transforms.py:1-10 states the program (temperature_then_stable_token_id_topk_inclusive_topp_original_max_minp_seeded_gumbel_v1) "intentionally makes no claim about generic SGLang or FlashInfer filter semantics". Inclusive top-p crossing, token-ID tie-breaking, and min-p relative to the original row max are exact-lane choices. If the generic path is going to score rollouts produced by a stock SGLang/FlashInfer sampler, that mismatch has to be settled explicitly, not left implicit. Related: @kiddyboots216's question on #60 about tests/ops/test_exact_sampling_transforms.py:83 (test_identity_top_p_one_is_full_support_despite_fp32_cumulative_overshoot) — the fp32 cumulative-overshoot handling at top_p == 1.0 should be resolved as part of pinning the shared program.

Proposed direction

1. Pin one sampling-transform program, backend-independent. Promote the exact_sampling_transforms program to the XoRL replay contract and document where it agrees/disagrees with stock SGLang and FlashInfer filters. Settle the top_p == 1.0 / fp32-cumulative question and the tie-break rule once, with a test that both paths run.
2. Extract one shared LM-head scoring op. Factor the common chunked-projection → fixed-order-LSE → support-mask → selected-logprob + VJP core out of the three implementations above; leave only the genuinely family-specific pieces (FP8/QLoRA weight access, TP topology) as injected policy.
3. Implement transforms on the generic modes. Support logprob_top_k/top_p/min_p and per-row logprob_temperature for compiled (and, where feasible, eager / quack_linear) so per_token_ce.py:274 and :337 stop being reachable for ordinary models. Chunked support-masking is compatible with the existing chunked CE structure.
4. Replace family gating with capability gating. Have generic models negotiate individual capabilities (exact sampling replay, per-row temperature, byte-exact PP) rather than needing membership in one of three hardcoded exact families. This is what collapses most of the 78 branch sites.
5. Fail loudly and in one place while (3) lands. Reject the request at admission with a single actionable message naming the unsupported ce_mode and the supported alternative — not two different late-stage errors from per_token_ce and model_runner, and never a silent unfiltered-support fallback.

Acceptance criteria

• A generic (non-exact-family) model can train on rollouts sampled with top_k/top_p/min_p and get logprobs on the filtered support, under ordinary body TP.
• One documented sampling-transform program, with the SGLang/FlashInfer divergences written down and tested.
• The chunked LM-head scoring core exists once; GLM-5.2, DSV4, and bi_fused consume it.
• No path silently scores against unfiltered support when the request carried transform metadata.
• Branch-site count on _exact_contract_family and friends materially reduced from 78.

Addresses #60.

_{References are against main @ 2b56f52.}

[kiddyboots216]

All of 1 / 2 / 3 / 4 / 5 sound good to me!