[feat] add DlrmHSTUOneRank (OneRank generative ranking)

[VERY-YCX]

Add DlrmHSTUOneRank: OneRank generative ranking with pluggable scorers

Summary

This PR adds DlrmHSTUOneRank (model tag 602), a OneRank-style generative ranking model built on top of the existing DlrmHSTU stack. It keeps the same input preprocessing, label decoding, point-wise losses, and metrics as DlrmHSTU, but replaces the shared FusionMTLTower head with per-task scoring over per-candidate, per-task representations produced by a structured token expansion inside HSTU.

Key changes

• Structured tokenization: each candidate is expanded into K interleaved (candidate copy, task token) pairs. The resulting task-private attention mask is encoded as two contiguous column intervals, so it fits the existing NFUNC=3 arbitrary-mask path without modifying the attention kernel.
• Situation Discernment (optional): per-task contextual query + softmax MHCA pooling over the request's own candidates. Disabled by default in our recommended config (see Deviations below).
• Cross-task head (optional): cascade-masked cross-task attention with strategic gradient detachment.
• List-wise InfoNCE loss: per-request softmax over the candidate list, jagged by candidate.sequence_length. No cross-rank gather is required.
• Pluggable scorer: dot_product (paper default), bilinear (identity-init per-task W_k), and mlp (per-task two-layer MLP). task_bias_init allows initializing per-task logits to logit(CTR_k).

Deviations from the OneRank paper

The implementation deliberately keeps SD and the cross-task head as config-gated optional modules, but our experiments on long-video TV recommendation found that:

• the rank-1 dot_product scorer tends to plateau at the global CTR prediction early in training;
• replacing SD + MHCA with simple mean-pooling and using the mlp scorer gives significantly better results in this scenario.

Therefore the recommended config for this use case is:

• situation_discernment unset,
• cross_task_head unset,
• scorer_type = ONERANK_SCORER_MLP,
• task_bias_init set to per-task empirical logit(CTR).

The other options are retained for ablation studies and for users who want to reproduce the original paper setup.

Backward compatibility

Only three upstream files receive subclass hooks (_build_output_modules, _build_stu_layer, _build_attn_func). The single behavioural change in stu.py is guarded by uses_arbitrary_mask and is False for all existing DlrmHSTU / UltraHSTU configs, so there is no impact on existing models.

Requirements

• kernel: CUTLASS or PYTORCH; TRITON is not supported because the NFUNC arbitrary-mask path is not implemented there.
• bf16/fp16 mixed precision is recommended for CUTLASS.
• Mid-stack attention truncation and KV-cache serving are not supported for this model.

Testing

• Unit tests for OneRankListwiseLoss covering numerical stability, masking edge cases, and gradient correctness.
• Model-level tests for DlrmHSTUOneRank on both PYTORCH and CUTLASS kernels, including kernel consistency checks and gradient reachability for task-token parameters.

[CLAassistant]

All committers have signed the CLA.

[github-actions]

Code review summary

Overall a well-structured addition: the structured-tokenization design is clearly motivated, the numerical handling in the list-wise loss is careful (per-segment max shift, exact clamps, masked-out requests, temperature clamp), and the tests are strong in places — the func-tensor vs. reference-mask test and the loss value+gradient reference checks are exactly the right assertions. The stu.py refactor was verified to be a no-op for existing SLA users (SLA func intervals are subsets of the fixed causal/target/contextual masks), and the model-level flip/length bookkeeping matches the DlrmHSTU baseline.

Main issue (see inline on onerank_tokenizer.py:230): the model as wired cannot be constructed with a default config — the stu.contextual_seq_len sentinel (-1) resolves to input_preprocessor.contextual_seq_len(), which is ≥ 1 for any contextual group, and OneRankSTULayer rejects anything but 0. This fires for the PR's own test configs too (they never set stu.contextual_seq_len: 0), so either the transducer needs to pin the value or the required config needs documenting and testing. The accompanying "baseline runs with contextual_seq_len = 0" claim is also inaccurate (the baseline gets a bidirectional contextual block).

Other noteworthy items (details inline):

• dlrm_hstu_onerank.py:249 — global-average-loss rescale uses total requests while the loss averages over valid requests only; biased cross-rank mean when mask ratios differ.
• onerank_jagged.py:39 — ~25 hidden device→host syncs per step from repeat_interleave without output_size=; all call sites have the total statically.
• dlrm_hstu_onerank_test.py — scorer_type variants / task_bias_init, cross-task gradient-detachment semantics, and listwise loss-input alignment have no coverage.
• onerank_sd.py:205 — per-task SD loop and cross-task pairwise projection are batchable/project-then-expandable.
• multi_task_rank.proto:209 — item_embedding_hidden_dim is silently ignored by this model; a few comment gaps noted there.
• Docs: ONERANK_OVERVIEW.md is referenced twice but does not exist in the repo, and there is no docs/source/models page / generative.rst entry / README row for the new model, unlike every comparable model (dlrm_hstu, ultra_hstu, hstu_match, sid_*).

Minor notes not worth inline comments: the comment in _build_output_modules claiming config_to_kwargs "does not translate enums or repeated fields" is inaccurate (json_format.MessageToDict translates enums to name strings and repeated fields to lists, and config_to_kwargs(onerank.cross_task_head) relies on exactly that for mask_type/hybrid_chain_task_names); nan_to_num(per_request, nan=0.0) in the loss does not absorb ±Inf (defaults rewrite them to the dtype's max finite value, which a valid request would carry into the sum); and incompatible HSTU options (attention truncation, KV-cache serving) fail at first forward rather than at construction, where the analogous constraints are already validated.

Note: the code-quality leg of this review was still running when this summary was posted; follow-up comments may follow.

[github-actions]

Follow-up: the code-quality leg of the review has now finished. It independently confirmed the two headline findings already posted (the stu.contextual_seq_len construction blocker at onerank_tokenizer.py:230, and the missing ONERANK_OVERVIEW.md / user-facing docs) and found no additional issues. It also verified — so maintainers need not re-litigate — that the NFUNC func-tensor intervals match the paper layout (including no-prefix and zero-candidate edges), the task-token slot selection in _compose_output is consistent with the tokenizer ordering, the double flip in predict restores both request and within-request candidate order, and the expand-based replicas sum gradients back into the candidate correctly. No further comments from this review.

[github-actions]

Code review — DlrmHSTUOneRank (automated, 5-area panel)

Overall this is a high-quality, well-hardened contribution: the NFUNC=3 mask-interval math checks out row-by-row (prefix/replica/task-token), the stu.py edits are verifiably a no-op for existing DlrmHSTU/UltraHSTU/SLA configs, the listwise InfoNCE is numerically robust (exact per-segment max-shift, denom >= 1 by construction, clamped temperature, all-masked batches tested finite), and the test suite goes well past the repo norm (CUTLASS↔PyTorch kernel consistency, gradient reachability into task tokens, detachment semantics).

11 inline comments posted. The ones I'd treat as blocking-or-explain:

1. The Chinese doc's example config cannot parse — task_configs uses non-existent fields (label_names/loss_type/metrics_set), and kernel: CUTLASS + input_dropout_ratio are nested inside stu {} where neither is a field. Since ModelConfig.kernel defaults to PYTORCH, a naive fix would silently run the reference kernel against the doc's own CUTLASS requirement.
2. model.proto oneof numbering — dlrm_hstu_onerank = 602 lands in the SID-generation block; multi-task rank models live in the 200s (dlrm_hstu = 205, ultra_hstu = 207). Oneof numbers are permanent; 208 is cheap now, expensive after release.
3. Cross-task Strategic Gradient Detachment — detaching after _k_proj/_v_proj means those weights only ever learn from diagonal entries, contradicting both the docstring ("projection weights still learn from every pair") and the "gradients are identical" comment. Forward values are identical, so no test catches it.
4. max_num_candidates is an unchecked data contract — a tail request exceeding it silently mis-autotunes every jagged kernel (eager) or fails mid-run inside a Triton op (compile). The tokenizer already pays the needed D→H sync; the check is free.
5. Test sampling gaps — sequence_timestamp_is_ascending=False may never be drawn in non-nightly CI (max_examples // 5, derandomized, 8 sampled axes), and the baseline needed a deterministic ordering test for exactly this contract; OneRank's listwise regrouping makes an order slip shape-invisible. SD has no value-level test; build_cross_task_mask has none at all.

Minor items (no inline comment each)

• pre-commit appears not run: README table isn't mdformat-normalized; onerank_head.py:215-218's magic trailing comma will be exploded by ruff-format; non-ASCII dashes in new comments outside tzrec/ops/.
• Comment inaccuracies: onerank_tokenizer.py:137-139 ("Negative on prefix rows" — torch.remainder takes the divisor's sign, so slot is always non-negative); dlrm_hstu_onerank.py:173-175 (config_to_kwargs does translate enums/repeated fields — see inline); proto says scorer_hidden_dim is "ONERANK_SCORER_MLP only" but validation is unconditional; task_bias_init's exact-length-K rule is undocumented while the doc example shows a single value for a multi-task model.
• ONERANK_OVERVIEW.md placement: an English user-facing explainer at repo root, outside the Sphinx tree and unreachable from any toctree, largely restating module docstrings; repo convention is Chinese user docs in docs/, implementation contracts beside the code. Consider folding section 5's plateau guidance into docs/source/models/dlrm_hstu_onerank.md and moving or dropping the root file (the proto and onerank_head.py cite it).
• Internal artifacts in permanent surfaces: the duration-bitmask enumeration {0,1,3,7,15,31}, "p99 = 10 on the target dataset", and "~100/900+ steps" run numbers in proto comments and docstrings don't generalize — state the model-level rationale instead.
• Loss edge cases: a zero-request batch on a rank gives 0/0 NaN that poisons all ranks through the DDP average — max(lengths.size(0), 1) completes the otherwise-thorough degenerate-path hardening (inherited from the point-wise path, but this module hardens everything else). Consider logits.float() under bf16 autocast, matching BCEWithLogitsLoss's fp32 promotion. logit_scale's gradient value is never compared against the reference (a .detach() on the scale would pass current tests).
• Fail-late config checks: TRITON kernel and attn_truncation_* are knowable at build time but only raise on first forward; the transducer already rejects five other misconfigs at construction.
• Perf micro-items: the stride-0 func-tensor view is re-.contiguous()'d per layer inside cutlass_hstu_mha (pre-existing for SLA, amplified by the inflated total_q); loss() adds one more scalar all-reduce over a value derivable from static shapes; _get_label re-decodes per listwise task; _preprocess returns a sum-of-maxes max_seq_len bound, looser than the max-of-sums the two D→H syncs were justified to avoid.
• Test consistency: pure-CPU classes tagged @mark_ci_scope("h20", "gpu") while sibling CPU classes are untagged; assertTrue(compute_metric()) only asserts a non-empty dict; OneRankHSTUTransducer's construction-time guards (interleaving, return_full_embeddings, contextual sentinel resolution) are untested; unused task_weight/output_dropout_ratio knobs in the test builder; _build_stu_layer/_preprocess overrides lack docstrings.

Areas reviewed: code quality, performance, test coverage, documentation accuracy, safety/validation — all five completed.

[github-actions]

Code review - DlrmHSTUOneRank (fresh 5-area panel: code quality, performance, test coverage, docs accuracy, security/validation)

All five areas completed against the current HEAD. This is an unusually careful PR, and it shows: we independently re-verified the pieces most likely to break silently and they hold up.

Verified sound (no action needed)

• stu.py hooks: the non-SLA path passes byte-identical kernel arguments, and the "no-op for SLA" claim checks out - build_sla_func_tensor folds contextual_seq_len (via effective_k2) and target isolation (via H_b) into the func intervals, so the func-visible set is a subset of both fixed masks and the previously-ANDed masks were redundant. UltraHSTU is untouched by the refactor.
• NFUNC=3 mask math row-by-row (prefix/replica/task-token intervals, no-prefix and zero-candidate edges, int32 plumbing), tokenizer slot order vs. _compose_output slot-1 extraction vs. task_configs order end-to-end.
• Flip bookkeeping in predict()/loss() (pre-flip num_targets, post-flip head inputs, flip-back of mt_preds) is order-consistent with jagged labels, and the enable_global_average_loss rescale (N_r / avg over a total-count denominator) is an unbiased global mean after DDP averaging.
• Loss hardening (detached per-segment max shift, exact denom.clamp(min=1), valid mask by multiply, nan_to_num, max(B,1)) prevents NaN/inf leakage into the all-reduce; the only collective is config-gated, so no rank-divergent participation.
• Registration surface: oneof 208 continues the 200s block, which_msg/RegisterABCMeta resolves the class, README/generative.rst/proto field docs are mutually consistent, and the arXiv link is the right paper.

New inline comments from this pass

• docs/.../dlrm_hstu_onerank.md:162 - mixed_precision: BF16 must be quoted (string field); the recommended CUTLASS example otherwise fails text-format parsing.
• docs/.../dlrm_hstu_onerank.md:179 - model export is listed as rejected, but nothing enforces it: the cand_seq_pk guard at export_util.py:147 covers only dlrm_hstu/ultra_hstu, and the export path traces the full forward without ever hitting cached_forward.
• onerank_jagged.py:71-75 - index_add_ accumulates in the input dtype; under the documented BF16 the recommended no-SD mean-pool builds z_k from bf16 atomics.
• onerank_tokenizer.py:551-556 - output postprocessor runs over the full 2K-row expansion before half the rows are dropped, and the non-contiguous slice pins the expanded buffer through backward.
• onerank_tokenizer.py:489-490 - both per-step D-to-H syncs are derivable from ints the preprocessor already synced (max_new_targets == max_targets * group_size exactly).
• dlrm_hstu_onerank_test.py:766,803 - the two CUTLASS tests (the only end-to-end NFUNC verification) inherit only the class-level "gpu" scope; repo precedent marks CUTLASS-wheel tests ("h20", "gpu").
• dlrm_hstu_onerank_test.py:395 - hypothesis GPU test lacks the sibling-test teardown_example/cleanup_cuda_memory().
• stu.py:536-547 - the SLA+CUTLASS behavioural no-op is untested through STULayer; a parity test would be cheap insurance.
• Two echo prior-panel comments and reinforce rather than add: SD per-task channel/stack order still has no value-level test (onerank_sd.py:211-220), and the class docstring still ties bf16/fp16 to both kernels (dlrm_hstu_onerank.py:69).

Prior-panel items we confirmed fixed at this HEAD: oneof 602 -> 208; docs moved into docs/source/models/ with rst/README entries; example-config field names; cross-task detachment now projects the detached input (pinned by gradient-equality tests); max_num_candidates runtime bound; output_size= on every repeat_interleave; zero-batch max(B,1) guard; logit_scale gradient compared against the reference; deterministic num_targets-order test.

Prior-panel items still open at this HEAD (no new inline; see the existing threads): temperature_init < 0.01 starts above the ln(100) clamp, so the learnable scale gets zero gradient forever (onerank_listwise_loss.py:67); the interleaving guard probes a private _enable_interleaving attribute fail-open (onerank_tokenizer.py:414); attn_truncation_* is the one incompatible knob rejected at first forward instead of construction (onerank_tokenizer.py:320); predict() duplicates the order-critical flip/assembly skeleton of the baseline instead of hooking it (dlrm_hstu_onerank.py:278); the example comment claiming FusionMTLTower requires mlp (md:101, same claim in dlrm_hstu_onerank_test.py:183).

Minor notes (no inline)

• The PR description says the recommended config leaves SD and cross_task_head unset, but the shipped doc example sets both - worth reconciling one way or the other.
• FP16 is documented for CUTLASS but only BF16 is tested; the listwise loss under low precision is asserted isfinite only (float64 parity is tested at unit level).
• New parameterized.expand tables omit name_func=parameterized_name_func (AGENTS.md); the loss test expands to opaque _0/_1/_2.
• Per-step aside already acknowledged in-code: the prefix split-expand-concat round trip copies the (long) prefix twice, and the prefix-timestamp concat is dead work since _compose_output discards the left half - architectural, fine to defer.

[VERY-YCX]

Housekeeping update: the 8 development commits (feature + four review-fix rounds + two follow-ups) have been squashed into the single head commit e31698d so the PR now reads as one complete change against master. The tree is byte-identical to the previously reviewed state plus two new fixes:

• 48bf972 — quote mixed_precision in the docs (it is a string field; the bare BF16 failed text_format.Parse)
• 39d056b — make the listwise-loss denominator FX-traceable (torch.fx.wrap leaf), fixing a TraceError under torchrec's train-pipeline rewrite that unit tests never exercise

All review feedback that was already addressed by the fix rounds (4f62fc9 / 90da9b5) has been replied to and resolved inline (22 threads). The remaining 17 open threads are the medium/low items from the last review round (doc-accuracy wording, minor test coverage, and optional perf suggestions) — they are tracked and will be handled in follow-up commits on this same branch.