DLRMv4: compute the last HSTU layer only for the rows the loss reads

[jiaweche]

Summary

In the DLRMv4 HSTU stack the loss reads one candidate row per sequence, so in the final
layer every other row's output is computed and then thrown away. This PR makes that last
layer compute U and Q for the candidate row alone against full-length K/V — K and
V stay full length, because the candidate still attends over the whole history. Nothing the
loss consumes changes.

On one MI350X at the production layer shape (512-dim transducer, 4 heads, qk and linear dim
128, bf16, batch 1024, 2.727 M tokens), the layer's forward+backward drops from
160.31 ms to 30.56 ms.

Opening as a draft: correctness and layer-level speedup are measured and reported below,
but a convergence run on yambda-5b is still outstanding. Details in
What is not yet validated.

Read First

The surviving row is the same function of the same inputs. Its reductions run in a different
order, so gradients are close rather than bitwise equal — the usual consequence of changing
tile shapes, not an approximation.

Dropout gets explicit care. The candidate row's mask is drawn using its own RNG row indices,
so the mask stream matches what the full-length path would have produced rather than merely
resembling it. stu_targets_only_test.py asserts the post-forward RNG state is exactly
equal between the two paths, alongside output and gradient agreement.

Gating

The rewrite is only valid under specific conditions, so it is checked rather than assumed,
and the failure mode is chosen per condition:

• Nine static conditions raise and name the blocker (kernel is not Triton, group norm on,
  not causal, not target-aware, windowed attention, contextual prefix, recompute flags off,
  inference). The knob is on for yambda-5b, so a config that can never satisfy it is a
  mistake worth surfacing — a silent fallback reads as "the optimization didn't help", and
  nobody re-checks a knob that appears to be working.
• Conditions that vary per batch fall back and log once per distinct reason: a batch that
  is not exactly one target per sequence, or eval when its own knob is off.
• Hardware falls back too, rather than raising. The indexed dropout mask exists only on
  the separated-RNG path (sm_100+ / MI350+). Below that it cannot be built, and no config
  change can fix it, so raising would turn a default-on knob into a hard failure on those
  GPUs. supports_indexed_output_dropout() is the single predicate for this, used by both the
  gate and the kernel.

Off by default in code; yambda_5b.gin turns it on for that benchmark. Both knobs are
overridable per-run by environment variable:

TRITON_HSTU_LAST_LAYER_TARGETS_ONLY=0   # disable (default in code)
TRITON_HSTU_TARGETS_ONLY_EVAL=0         # keep eval on the full-length path

FLOPs accounting had to change

The counter charged every layer at full length. Left alone it would have inflated HFU by
roughly the size of the saving, exactly when the work was being removed. The metric is now
split:

• fill keeps its original meaning — ragged tokens over padded budget.
• exec is new, and carries the targets-only discount. HFU divides by exec.
• MFU deliberately keeps the dense yardstick. It therefore rises when this is enabled,
  and must not be read as improved utilization. Use HFU or ms/step for that.

Evidence

Parity — modules/tests/stu_targets_only_test.py, 6 tests, all pass on one MI350X. The
load-bearing one checks forward output at the candidate row, every parameter gradient, and
exact RNG-stream alignment against the full path. The hardware-decline test asserts by
substitution, since that condition cannot be reached on hardware that has the separated-RNG
path.

No regressions — layer_norm_test, fake_signature_test, hstu_compute_test, and
stu_test give 5 failed, 10 passed on this branch and 5 failed, 10 passed on pristine
master: same five tests, same counts. Both pre-existing failures are unrelated to this
change (stu_test hits a missing torch.ops.fbgemm.asynchronous_complete_cumsum;
hstu_compute_test fails a float32 tolerance on a training=False example, a path this PR
does not touch).

Layer speedup — one MI350X, bf16, production layer config. The full arm backpropagates
from the candidate rows only, so both arms produce exactly the rows the loss reads; the full
arm's forward and backward still run at full length, which is the cost being removed.

batch	length	tokens	full	targets-only	speedup
256	2663 uniform	0.682 M	41.26 ms	13.64 ms	3.03×
1024	2663 uniform	2.727 M	160.31 ms	30.56 ms	5.25×
1024	2663 ±50% ragged	2.759 M	186.94 ms	27.02 ms	6.92×
1024	4086 uniform	4.184 M	333.64 ms	41.42 ms	8.06×

The gain grows with sequence length because the removed work is the O(L²) attention term
while what remains is linear in tokens. At equal token count the ragged case costs the full
arm 17% more than uniform (sum of L² rises with length variance) while the targets-only arm
does not move.

End-to-end in the trainer — the debug synthetic dataset exercises the transducer wiring,
the compact postprocess, and the accounting. All three arms train to completion (rc=0):

arm	gate	reported
10 targets/seq	declines, names the reason	fill=87.3% exec=87.3%
1 target/seq, knob on	engages	fill=81.3% exec=65.9%
1 target/seq, knob off	not consulted	fill=83.7% exec=83.7%

exec drops below fill only when the rewrite actually runs, so the discount cannot be
claimed spuriously.

What is not yet validated

Everything above is single-GPU, and either synthetic or layer-level. Two things are still
open, both needing the real yambda-5b on 8 GPUs:

1. End-to-end step time at production scale. The layer measurement bounds the gain; the
   last layer's share of total step time sets the actual number.
2. Convergence. That seeds still land inside the RCP range. This is the one that matters
   for a reference change, and it is why this is a draft.

I am running these now and will post results here. Review of the approach and the gating
design is welcome in the meantime.

Provenance

Ported from an AMD fork where this has been running as part of MI350X/MI355X benchmarking.
A companion optimization in that fork (FAST_INTERIOR, an interior-tile masking
specialization) is deliberately excluded here: it measured +0.02% end-to-end, which does
not justify its complexity in a reference implementation.

One difference from the fork is worth flagging for anyone comparing numbers. The fork reports
118.01 → 27.35 ms at the same 2.727 M tokens (4.31×). The targets-only arms agree closely
(30.56 vs 27.35 ms) since that path is Triton in both trees, but the full arms do not
(160.31 vs 118.01 ms) because the fork's baseline uses its own hand-scheduled backward. The
ratio here is larger than the fork's for an uninteresting reason: this baseline is slower, not
this optimized path faster.

[nanz-nv]

Thanks for proposing this optimization!

I think it is legitimate and mathematically equivalent. It skips unnecessary computation for activations that would be discarded anyway.

One minor note: I noticed a small tweak to the dropout implementation intended to preserve the mask as if dropout were applied to the full activation. This would be useful if the other skipped operations could also be made bit-matching, so that the entire optimization would be bit-matching. Otherwise, I do not see this as particularly useful. Under the MLPerf rules, optimizations should preserve high-level mathematical equivalence; bit matching is not required.

[ShriyaRishab]

WG meeting discussion 9/3/26:

1. Should we merge this into the reference since we're very close to submissions and we shouldn't increase burden? Let's merge since its better for the community and it doesn't change the math/algorithmic behavior
2. Can @pavanky and Linjin to review this optimization?
3. Please fix CLA

[LinjianMa]

Looks good to me

[ShriyaRishab]

@jiaweche can you please fix the CLA?