[feat] ULTRA-HSTU Mixture of Transducers (MoT)

[tiankongdeguiji]

Summary

Adds UltraHSTU as a new top-level model type implementing the four ULTRA-HSTU optimizations from Ding et al., 2026 (arXiv:2602.16986) on top of DlrmHSTU:

• Semi-Local Attention (SLA) — HSTU.stu.sla_k1 (local window) + sla_k2 (global prefix); requires Kernel.CUTLASS or Kernel.PYTORCH.
• Mid-stack Attention Truncation — HSTU.attn_truncation_split_layer + attn_truncation_tail_len; drops UIH prefix after the split layer, preserving contextual + targets.
• Mixture of Transducers (MoT) — repeated HSTU hstu in UltraHSTU runs N parallel transducers, one per UIH channel; per-candidate outputs concatenate on the embedding dim before the multi-task tower. Candidate and contextual groups stay shared. Channel selection is via a single new optional string HSTU.name field that replaces the default uih prefix on UIH-side feature_group keys (e.g. name="uih_click" reads uih_click.sequence / uih_click_action.sequence / ...). Empty name preserves all current DlrmHSTU configs.
• Selective rematerialization — already on STU via recompute_normed_x_in_backward / recompute_uvqk_in_backward; surfaced + documented for ULTRA-HSTU.

Single-channel UltraHSTU (one hstu entry, empty name) is behaviorally equivalent to DlrmHSTU; SLA and Attention Truncation are independently usable in that mode too.

Implementation notes

• Proto: one optional string name on HSTU; new UltraHSTU model message in multi_task_rank.proto mirroring DlrmHSTU's top-level fields with repeated HSTU hstu; UltraHSTU ultra_hstu = 207 in model.proto's oneof.
• Runtime: UltraHSTU subclasses DlrmHSTU via two small extracted hooks (_build_transducer, _stu_embedding_dim); the wrapper _HSTUTransducerStack is bypassed for single-channel configs so AOTI / JIT trace shapes are unchanged. STULayer exposes attn_func_static_sig (a colon-separated :-string of init-time SLA constants) as a @property; STUStack precomputes a per-layer prev_attn_func_sig list at construction so the forward loop has zero sig bookkeeping and is fx-trace-safe by construction.
• SLA + truncation infra hardened for fx symbolic-trace + AOT-export end-to-end (the cutlass kuairand integration test was the first to exercise them under export). Fixes:
  • build_sla_func_tensor: unconditional int32 cast (no Proxy.dtype check); torch.diff + repeat_interleave instead of searchsorted on a slice (avoids Inductor SliceView.get_stride() NotImplementedError).
  • compute_stu_truncation_plan returns precomputed total-int fields on STUTruncationPlan (total_dropped / total_kept / total_prefix / total_rest) so apply_stu_truncation_plan can pass them as total_len_* into split_2D_jagged and skip its .item() fallback under FakeTensor.
• Tests + config + docs:
  • Renamed tzrec/tests/configs/dlrm_ultra_hstu_cutlass_kuairand_1k.config → ultra_hstu_cutlass_kuairand_1k.config; converted to ultra_hstu with two channels (uih_click, uih_view); added test_rank_ultra_hstu_cutlass_train_eval_export in rank_integration_test.py.
  • New kuairand-mot-1k parquet (split UIH by action_weight bits at preprocess time) wired into scripts/ci/ci_data.sh (oss-accelerate); existing kuairand-1k URLs also moved to oss-accelerate.
  • tzrec/models/ultra_hstu_test.py parametrizes over (single, multi_uniform, multi_hetero) channel topologies × NORMAL / FX_TRACE / JIT_SCRIPT / AOT_INDUCTOR × PYTORCH / TRITON.
  • tzrec/ops/hstu_attention_utils_test.py parametrizes every numerical-correctness test over NORMAL + FX_TRACE for direct fx-trace regression coverage.
  • docs/source/models/ultra_hstu.md (NEW) + README + generative.rst toctree.
  • tzrec version bump 1.1.16 → 1.1.17.

Test plan

• python -m tzrec.modules.gr.preprocessors_test — back-compat (empty name defaults to existing keys).
• python -m tzrec.modules.gr.hstu_transducer_test — back-compat.
• python -m tzrec.modules.gr.stu_test — SLA cache + truncation back-compat.
• python -m tzrec.modules.gr.preprocessors_test — back-compat.
• python -m tzrec.ops.hstu_attention_utils_test — 27 parametrized cases (eager + fx-trace).
• python -m tzrec.models.dlrm_hstu_test — DlrmHSTU refactor is behavior-preserving (full hypothesis sweep).
• python -m tzrec.models.ultra_hstu_test — 3 channel topologies × hypothesis sweep across all graph types and kernels.
• python -m unittest tzrec.tests.rank_integration_test.RankIntegrationTest.test_rank_ultra_hstu_cutlass_train_eval_export — end-to-end train → eval → export → predict on kuairand-mot-1k with two MoT channels under CUTLASS + AOT export.

🤖 Generated with Claude Code

[github-actions]

Renumbering PEPNet from 206 → 207 and reusing tag 206 for UltraHSTU is wire-format-breaking. Any binary-serialized ModelConfig that previously stored pepnet at tag 206 will now silently parse as ultra_hstu (or be discarded as a type mismatch). In-tree configs are loaded via text_format/json_format (name-keyed) so the in-repo path is unaffected, but this is bad proto hygiene and dangerous for any downstream tooling, exported pipeline.config bundles, or external consumers.

Suggested fix: keep PEPNet pepnet = 206 at its original tag and assign UltraHSTU a fresh, never-published tag (e.g. 208). Tag numbers should be treated as immutable once published.

[github-actions]

Two related validation gaps in this block:

1. Single-channel mode skips all name validation. A user with one channel can set name="candidate" (or any reserved feature-group name); _build_transducer will then call embedding_group.group_total_dim("candidate") — which exists, since candidate is the shared candidate-side group — and the UIH path will silently consume candidate-side embeddings without raising. Worth checking that single-channel name (when non-empty) is also non-reserved.

2. No charset/whitespace check on name. The proto field is an unconstrained optional string. A name containing . would produce keys like "foo.bar.sequence"; EmbeddingGroupImpl.group_total_dim splits on . and would resolve via "foo", mapping the user's intent to an unrelated group. Whitespace or look-alike unicode would mismatch the dict key from the feature-group config, producing a confusing KeyError deep in forward.

Suggested: validate name once (non-reserved, no ., matches ^[A-Za-z_][A-Za-z0-9_]*$) for every channel regardless of count.

[github-actions]

Two small concerns on the stack construction path:

1. Silent friendliness gap on missing groups. If a channel sets name="uih_click" but the user forgets to define one of the four required groups (uih_click, uih_click_action, uih_click_watchtime, uih_click_timestamp), the missing UIH-group surfaces here as a bare KeyError from group_total_dim, while a missing _action/_watchtime/_timestamp group only surfaces at runtime inside ContextualInterleavePreprocessor.forward. Consider asserting the four groups exist up-front (skipping _action/_watchtime when those encoders aren't configured) so misconfigs fail with a clear message at construction time.

2. _HSTUTransducerStack always discards the second tuple element. The stack hardcodes return_full_embeddings=False per channel and returns None as the second slot, but if a future change ever flips return_full_embeddings=True on a sub-transducer (e.g. for an aux-loss head), cand[1] is silently dropped. Worth a one-line assert not transducer._return_full_embeddings in the stack constructor and a note in its docstring.

[github-actions]

Two coverage gaps from hardcoding _CHANNELS = ("a", "b") and embedding_dim=512 for every channel:

1. Single-channel path is never exercised. The branch at ultra_hstu.py:118-119 (if len(transducers) == 1: return transducers[0]) — the documented "behaves like DlrmHSTU" mode that bypasses _HSTUTransducerStack entirely — is the easiest path to silently break in a future refactor (e.g. shape regression: single tensor vs. concatenated). Recommend a separate test case (or parametrize over (("a",), ("a","b"))) and assert isinstance(model._hstu_transducer, HSTUTransducer) for the single-channel case.

2. Heterogeneous channel dims are not exercised. With both channels at embedding_dim=512, sum(c.stu.embedding_dim for c in channels) is mathematically indistinguishable from N * dim or max(...) — a buggy _stu_embedding_dim would still pass. Use distinct dims (e.g. 256 and 512) in at least one case so the sum is the only correct answer.

Also: there are no negative tests for the __init__ validation assertions (len(channels) >= 1; >= 2 channels need unique non-empty names) — easy to add and they protect user-facing error paths.

[github-actions]

Per-channel embedding_name=f"{channel}_..._emb" gives every channel its own embedding table for the same physical feature (video_id, video_cat). Users will copy this pattern, and in production it multiplies sparse-parameter count, TBE forward/backward work, and all-to-all communication volume by N — typically the dominant cost in a sharded recsys deployment. The integration config (dlrm_ultra_hstu_cutlass_kuairand_1k.config) implicitly shares table names because both channels point at the same uih_seq__video_id, so the test config and the integration config disagree on the recommended pattern.

Suggest making the test share embedding_name across channels by default to set the right example, and add a class/proto docstring note recommending shared embedding_name unless per-channel tables are explicitly desired.

[github-actions]

Fields 2-7 of UltraHSTU are an exact copy of DlrmHSTU fields 2-7; only field 1 (hstu repeated vs single) differs. The Python _init() already references all of these via self._model_config.X and works for both messages purely by field-name duck typing — that contract silently breaks the day someone adds a field to one and forgets the other.

Lower-effort options: add a // keep in sync with DlrmHSTU 2-7 TODO comment to both messages. Higher-effort: extract a shared sub-message embedded in both, or just make DlrmHSTU.hstu repeated and retire UltraHSTU entirely.

[github-actions]

Review Summary

Solid, well-scoped addition. The _init / _build_transducer / _stu_embedding_dim extraction in DlrmHSTU is a clean Template Method seam, the MoT key-routing in ContextualInterleavePreprocessor is small and self-contained, and the test matrix covers NORMAL / FX_TRACE / JIT_SCRIPT / AOT_INDUCTOR × PYTORCH / TRITON. Most findings are robustness / contract hygiene, not correctness.

Highest-impact items (see inline comments for details):

• tzrec/protos/model.proto:73 — renumbering PEPNet 206→207 to give UltraHSTU tag 206 is wire-format-breaking. Reusing a previously-published tag for a different message is dangerous for any downstream binary-serialized ModelConfig. Prefer keeping PEPNet=206 and giving UltraHSTU a fresh tag (e.g. 208).
• tzrec/models/ultra_hstu.py channel-name validation — single-channel mode skips name validation entirely; name has no charset/whitespace/reserved-token check, so values like "candidate" or "foo.bar" silently misroute via the embedding-group .split('.') lookup.
• _HSTUTransducerStack (ultra_hstu.py:41-48) — silently discards the second tuple element from each sub-transducer; one-line invariant assertion would prevent a confusing future regression. Also worth asserting up-front that all four required UIH groups exist for each channel, instead of letting the failure surface deep in forward.
• Test coverage gaps (ultra_hstu_test.py) — single-channel degradation path, heterogeneous per-channel embedding_dim, and __init__ validation negative tests are not exercised. The hard-coded per-channel embedding_name in the test also models a pattern that would multiply embedding tables N-fold in production; consider sharing embedding_name across channels to match the integration config.
• multi_task_rank.proto — UltraHSTU fields 2-7 mirror DlrmHSTU exactly; consider a sync-comment or refactor to prevent silent drift.

Performance note (informational, not blocking): _HSTUTransducerStack.forward runs the N transducers in a serial Python for and re-executes the candidate-side preprocessor preamble (including N fx_int_item D2H syncs) per channel. Fine for 2 channels, but the design clearly aims at higher N where this becomes the dominant cost. Worth profiling and potentially hoisting the candidate-side state out of the inner loop or using CUDA streams.

Docs: the deferred docs/source/models/ultra_hstu.md is acknowledged in the PR body. Worth also touching dlrm_hstu.md group-name section and generative.rst / README.md model table when that doc lands.

[github-actions]

The proto field names cited here don't exist. The actual fields in tzrec/protos/module.proto:249,251 are recompute_normed_x and recompute_uvqk (no _in_backward suffix — that suffix only appears as the internal kwarg in tzrec/ops/hstu_compute.py). A user copy-pasting recompute_normed_x_in_backward: true into a stu { ... } config will get a proto-parse error. Please rename to the actual proto field names.

[github-actions]

Two related validation gaps worth closing up-front:

1. Reserved-name collisions are silently accepted. name="uih", name="candidate", or name="contextual" on a multi-channel config will pass the all(names) and uniqueness check, but the derived UIH-side keys (preprocessors.py:151-154) will then alias the candidate / contextual / default-uih groups. Suggest rejecting reserved names and *_action / *_watchtime / *_timestamp suffixes that would collide with the per-channel naming convention.

2. Missing feature_group is reported deep, not at construction. A typo in name surfaces as a KeyError in preprocessors.py:373 (grouped_features[f"{self._uih_key}.sequence"]) at forward time, or as whatever embedding_group.group_total_dim(uih_group) raises mid-_build_transducer. Validating up-front that for every channel the four groups (<name>, <name>_action, <name>_watchtime, <name>_timestamp) exist on self.embedding_group and listing the missing keys would make misconfigurations obvious.

[github-actions]

Coverage gaps worth closing:

1. Three new __init__ error branches are untested (ultra_hstu.py:84-90): zero channels, len>=2 with empty name, len>=2 with duplicate names. All are CPU-only assertion checks — trivial to add and would prevent a regression silently dropping the validations.

2. GPU-skip hides CPU-only structural assertions. Lines 396-403 (isinstance(..., _HSTUTransducerStack) and _stu_embedding_dim() == sum(d for _, d in channel_specs)) are pure-Python checks that gate the central MoT structural invariant (the only-correct-answer for multi_hetero is sum, not max or N*dim). Hiding them behind @unittest.skipIf(*gpu_unavailable) means CPU CI never validates them. Recommend splitting into a CPU-only structural test plus the GPU forward test.

3. No single-channel-with-empty-name case — the ("single", [("a", 256)]) parametrize case uses a non-empty name, so the _build_transducer fallback to the "uih" group at ultra_hstu.py:98 is never exercised. This is the documented "无痛迁移 from DlrmHSTU" path.

4. Output is shape-checked only, not value-checked. A two-channel parity test (identical seeds for both channels, then assert the two halves of torch.cat(..., dim=-1) along the embedding dim are equal) would catch a broken concat order or wrong index in _HSTUTransducerStack.forward for negligible cost.

[github-actions]

N transducers run serially in a Python for loop on a single CUDA stream — each transducer launches its own preprocessor + STU stack + postprocessor + many small jagged kernels. With per-channel embedding_dim typically 32–128, kernel-launch latency and SM under-occupancy will dominate at small batch sizes, so MoT scaling won't be linear in N.

The user-facing doc at docs/source/models/ultra_hstu.md:11 calls this "模型并行运行 N 个 HSTUTransducer", which oversells the GPU behavior — eager + AOT-export both serialize today. At minimum, please soften the doc claim. Optionally, on the eager-train path you could record each iteration on a dedicated torch.cuda.Stream and join before torch.cat (skip on AOT-export).

Also, no shape/dim sanity check before torch.cat(..., dim=-1) — divergent per-candidate row counts across transducers (e.g. a future per-channel truncation behavior change) would silently misalign rows. A cheap precondition assert would prevent a hard-to-debug regression.

[github-actions]

The cache key omits _max_attn_len and _causal, both of which affect NFUNC mask semantics. It happens to be safe today because STUStack constructs every layer in a stack with identical kwargs (hstu_transducer.py:111), so all layers in a stack share the same (max_attn_len, causal). But two layers that differ only in max_attn_len (a plausible future heterogeneous-layer config) would silently reuse a stale mask via this sig.

Defensive fix — include both in the sig:

return (
    f"{self._sla_k1}:{self._sla_k2}:{self._contextual_seq_len}:"
    f"{self._max_attn_len}:{int(self._causal)}:"
    f"{self._num_heads}:{int(self._target_aware)}"
)

[github-actions]

Code Review Summary

Five reviewers covered code quality, performance, tests, docs, and security. Posted 5 inline comments on the highest-signal findings; collecting the rest here.

Most actionable (see inline)

• Doc bug: docs/source/models/ultra_hstu.md:13 cites proto fields recompute_normed_x_in_backward / recompute_uvqk_in_backward — these don't exist; actual proto fields are recompute_normed_x / recompute_uvqk. A copy-paste from the docs into a config will fail to parse.
• Validation: UltraHSTU.__init__ accepts reserved channel names (uih, candidate, contextual) and defers feature-group existence to a deep KeyError.
• Tests: Three new error branches in UltraHSTU.__init__ are untested; structural MoT assertions in ultra_hstu_test.py:396-403 are CPU-only checks gated behind gpu_unavailable.
• Docs vs reality: "模型并行运行 N 个 HSTUTransducer" oversells GPU scaling — _HSTUTransducerStack runs serially on a single CUDA stream.
• SLA cache key: STULayer.attn_func_static_sig omits max_attn_len and causal (safe today since all layers in a stack share kwargs, but a foot-gun for future heterogeneous stacks).

Other notes worth scanning

• arXiv ID 2602.16986 — plausible but the high .16986 suffix for a Feb 2026 submission is on the edge of monthly volume. Worth double-checking the link before publishing.
• scripts/ci/ci_data.sh — no set -euo pipefail and no checksum verification on the new kuairand-mot-1k parquet. The 32-hex hash in the URL is purely advisory; a partial download will silently produce CI noise.
• MoT redundant work — under MoT the _combine_embeddings interleave mask (preprocessors.py:255-277) and contextual-prefix concat_2D_jagged (preprocessors.py:301-324) re-compute identical results across channels. Could share across channels if a profile shows them hot.
• No MoT correctness parity test — current end-to-end test asserts only output shapes. Two channels with identical seeds vs. single-channel double-concat would catch a broken torch.cat(dim=-1) order or wrong index in _HSTUTransducerStack.forward.
• STU.truncate_input declared @abc.abstractmethod with a docstring asking subclasses to "override and raise NotImplementedError" — that contradicts the abstract decorator. Either drop the decorator (provide a default raise body on the base) or drop the docstring guidance.
• Selective rematerialization defaults — the new docs claim defaults are true, which matches module.proto:249,251. Just calling out that this is a behavior-affecting default users will likely not notice.

Praise

• The fx-trace + AOT-export hardening is well-thought-out: STUTruncationPlan precomputes total_dropped / total_kept so split_2D_jagged skips its .item() fallback under FakeTensor; STUStack precomputes _prev_attn_func_sig_per_layer so the forward loop has zero sig bookkeeping. The _HSTUTransducerStack bypass when len == 1 is a nice touch — keeps single-channel UltraHSTU traceable identically to DlrmHSTU.
• Good defensive choice in build_sla_func_tensor: the torch.diff + repeat_interleave rewrite (avoiding a slice-SliceView Inductor crash) and the clamp(L - T, min=0) to dodge silent NaN attention.
• Test parametrization across (single, multi_uniform, multi_hetero) × all four TestGraphTypes is thorough on the structural side; the suggestions above are about what's not yet covered, not the breadth.