[bugfix] thread contextual_seq_len from preprocessor to STULayer (proto sentinel + truncation total_uih_len + AOTI-friendly SLA builder)

[tiankongdeguiji]

Summary

PR #450 added if "contextual_seq_len" not in stu: to HSTUTransducer.__init__ intending to forward _input_preprocessor.contextual_seq_len() into every STULayer when the user hadn't pinned stu.contextual_seq_len. That guard has been silently unreachable since day 1 because tzrec/utils/config_util.py runs MessageToDict(..., including_default_value_fields=True, ...) and the proto field was optional uint32 with no explicit default — so the dict always contains contextual_seq_len: 0 and the membership test never matches. Every STULayer was constructed with contextual_seq_len = 0 regardless of the preprocessor's value.

For models without SLA + attn_truncation + multi-channel MoT (e.g., dlrm_hstu_cutlass) the wrong value was numerically harmless. For ultra_hstu_cutlass_kuairand_1k (sla_k1=256, sla_k2=32, attn_truncation_split_layer=2, attn_truncation_tail_len=512, attn_num_layers=4, 2 MoT channels, 6-feature contextual group), it manifests as l2_loss_watchtime: 1.92e9 on step 0 followed by a forward-pass CUDA IMA on step 1 (see PR #500 CI run 25497698553 / job 74821854417, and the matching Hopper repro on PR #501 H20 CI run 25543727919 / job 74975071509).

This PR closes the loop end-to-end:

1. Proto sentinel (tzrec/protos/module.proto): optional uint32 contextual_seq_len = 13; → optional int32 contextual_seq_len = 13 [default = -1];. Mirrors the scaling_seqlen pattern from PR [feat] stu.scaling_seqlen + drop autotune assert strip #500.

2. Sentinel-style guard (tzrec/modules/gr/hstu_transducer.py): replace the unreachable if "contextual_seq_len" not in stu: with if stu.get("contextual_seq_len", -1) < 0: stu["contextual_seq_len"] = self._input_preprocessor.contextual_seq_len().

3. Post-truncation total_uih_len fix (tzrec/modules/gr/hstu_transducer.py:_replay_truncation_state): change plan.total_kept - total_targets to plan.total_kept + plan.total_prefix - total_targets. total_kept is the post-truncation "rest" portion only (excluding contextual prefix); _postprocess builds uih_offsets from seq_lengths - num_targets which still includes the prefix per sample, so the under-count made the Triton/CUTLASS split_2D_jagged allocate OutA smaller than offsets_left[-1] and write OOB on the very first forward step. (The unit test was asserting the buggy expected value too — also corrected.)

4. AOTI-friendly SLA NFUNC builder (tzrec/ops/hstu_attention_utils.py:build_sla_func_tensor): with cs > 0 actually flowing through, AOTI export hit two new Inductor-compile failures the master cs = 0 path never tripped:

   • unsqueeze(0).expand(nheads, 3, total_q).contiguous() triggers Inductor's combine_contiguous_dims for a 3-D iteration with symbolic last dim → ModularIndexing(idx, total_q, 3) → sympy ZeroDivisionError.
   • seq_offsets[batch_ids] / seq_lengths[batch_ids] / num_targets[batch_ids] (where batch_ids = repeat_interleave(arange(B), seq_lengths)) materializes batch_ids as a kernel input and emits tl.device_assert(0 <= batch_ids < B). Under AOTI compile-time autotune the bench fills the input with rand_strided int32 garbage that never satisfies the bound — assert fires, export aborts with cudaErrorAssert.

   Fixes (no Python custom ops, no torch._inductor.config overrides):

   • Defer .contiguous() to the consumer: return func_2d.unsqueeze(0).expand(nheads, 3, total_q). cutlass_hstu_mha is @torch.fx.wrap'd (FX leaf) and already calls .contiguous() itself (tzrec/ops/_cuda/cutlass_hstu_attention.py:136); deferring keeps combine_contiguous_dims outside Inductor's compile boundary.
   • Replace indirect indexing with torch.repeat_interleave(values, lengths) directly per slot (seq_offsets_per_pos = repeat_interleave(seq_offsets[:B].contiguous(), seq_lengths), same for L / T). Mathematically equivalent to values[batch_ids], but Inductor's lowering for the (values, lengths) overload is a single cumsum-driven scatter/gather that doesn't surface a separate batch_ids kernel input — no assert_indirect_indexing codegen, no autotune-bench bounds-check failure.
   • compute_stu_truncation_plan: new_x_offsets = x_offsets - offsets_head (mathematically equivalent to offsets_prefix + offsets_tail, but no arange(B+1)*cs chain feeding into the SLA builder's seq_offsets — keeps cs (and via it B) out of the consumer's SymInt graph).

Wire-compat

• A uint32=0 from any old serialized config decodes as int32=0. The < 0 sentinel (not <= 0) preserves the legacy semantic that an explicitly-set 0 means "no contextual" rather than "use preprocessor".
• tzrec/protos/module_pb2.py is gitignored and regenerated by scripts/gen_proto.sh at install time — no checked-in .py change. Regenerated descriptor sanity-checked locally: d['contextual_seq_len'] == -1 for an empty STU under including_default_value_fields=True.
• AOTI exported .pt2 references no tzrec:: Python custom ops — runs in pure C++ runtime.

Test plan

• Local python -m unittest tzrec.modules.gr.hstu_transducer_test — 7/7 pass (covers ctx in {0, 3} × (split, tail) in {(0,0), (1,4)} × interleave matrix).
• Local python -m unittest tzrec.ops.hstu_attention_utils_test — 27/27 pass (truncation-plan parity).
• Local python -m unittest tzrec.ops.hstu_attention_test.HSTUAttentionTest.test_sla_attn_cutlass — 60 hypothesis examples, OK (CUTLASS NFUNC parity vs PyTorch reference unchanged).
• Local python -m unittest tzrec.tests.rank_integration_test.RankIntegrationTest.test_rank_ultra_hstu_cutlass_train_eval_export — Ran 1 test in 162.842s, OK (full train_eval + eval + export + predict on 4×A10 + cu129 + fbgemm_gpu_hstu wheel).
• Local proto round-trip: empty STU → {'contextual_seq_len': -1, …} → sentinel < 0 triggers preprocessor fallback; explicit 0 → respected as 0; explicit positive → respected verbatim.
• CI Hopper lane runs test_rank_ultra_hstu_cutlass_train_eval_export and it passes (was the originally failing test).
• CI confirms test_rank_dlrm_hstu_cutlass_train_eval_export still green (different code path, regression spot-check).

🤖 Generated with Claude Code

[github-actions]

This AOTI-compile-time-autotune workaround also runs on the eager runtime forward path (STULayer.forward → build_sla_func_tensor), so every CUTLASS layer of every batch pays it.

Old form: 1 repeat_interleave(arange(B), seq_lengths) (≈1 cumsum + 1 scatter, length total_q) + 3 cheap B-element gathers.
New form: 3 repeat_interleave(values, seq_lengths) calls (each ≈1 cumsum + 1 scatter at length total_q) + the narrow(...).contiguous() copy.

Net: roughly +2 cumsum kernels and +2 full-length scatters per SLA layer per forward. With a stack of N CUTLASS layers, that's 4N extra kernel launches per forward.

Consider gating the new form behind an export-only path (e.g., torch.fx._symbolic_trace.is_fx_tracing() / a module-level _export_mode flag) and keeping the original batch_ids = repeat_interleave(arange(B), seq_lengths) + 3 gathers on eager — the autotune device-assert only fires under AOTI compile-time bench, not at eager runtime.

Side note on line 113: seq_offsets_i32.narrow(0, 0, B) on a 1-D contiguous source is already contiguous, so .contiguous() is a no-op. If it's there to defeat an Inductor SliceView issue analogous to the seq_offsets[1:] problem documented at line 92, please add a one-line comment so a future cleanup pass doesn't drop it.

[github-actions]

Code review summary

Reviewed via code-quality, performance, test-coverage, doc-accuracy, and security subagents. The fix is well-scoped and the inline comments justifying each AOTI workaround are excellent. Three inline comments left covering the main concerns.

Headline findings

• Perf: the new 3× repeat_interleave(values, lengths) form in build_sla_func_tensor is an AOTI compile-time-autotune workaround but also runs on every layer of every eager forward (~4 extra kernels × N CUTLASS layers / forward). Consider gating behind an export-only path. (inline)
• Test coverage: the headline bug fix (stu.get("contextual_seq_len", -1) < 0) isn't actually pinned by any test — every existing case has the key absent from the stu dict, so the old not in form would pass identically. A 3-row parameterization on the stu-dict key would pin the contract. (inline)
• Doc: STUTruncationPlan.total_kept's field doc doesn't mention it excludes the contextual prefix, even though the new _replay_truncation_state + plan.total_prefix add-back depends on that. (inline)

Smaller items (skipped from inline)

• hstu_transducer.py:108 comment says "default-zero fill" but the proto now fills -1, not 0 — minor wording cleanup.
• The proto comments for contextual_seq_len (line 256-257) and scaling_seqlen (line 263) both use the < 0 (default) = sentinel pattern but with slightly different phrasing; standardizing helps grep-ability.
• Mathematical equivalence of new_x_offsets = x_offsets - offsets_head vs. the old offsets_prefix + offsets_tail confirmed; the rerouting rationale (keeping cs out of the seq_offsets dependency chain for AOTI) is correctly captured in the inline comment at lines 258-270.

Security

No noteworthy security/safety findings. The proto-sentinel-then-validate layering correctly prevents any negative contextual_seq_len from reaching the kernels via the config path; the in-kernel < 0 validators in build_sla_func_tensor and compute_stu_truncation_plan retain defense-in-depth. Wire-compat claim (uint32→int32) holds for any value an old client could realistically encode.

Positive

• Root cause analysis in the PR body is precise (the including_default_value_fields=True interaction is exactly what made the previous not in guard unreachable).
• The + plan.total_prefix add-back fix in _replay_truncation_state is correctly justified by the layout mismatch between total_kept (rest-only) and _postprocess's seq_lengths - num_targets (includes prefix).
• Removing the _sla_broadcast_func_to_heads custom op in favor of expand + caller-side .contiguous() is a wash perf-wise and reduces export-side complexity.