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

recommendation/generative_recommenders/dlrm_v4/train/_env_bootstrap.py

@@ -19,9 +19,28 @@
 logger: logging.Logger = logging.getLogger(__name__)
 
 
+def _bind_bool(name: str, value: Optional[bool]) -> None:
+    """Push a gin-bound boolean into the environment; a pre-set env var wins.
+
+    Same env-wins rule as the binding written out below, factored out because
+    each additional kernel flag needs it identically.
+    """
+    if name in os.environ:
+        logger.info(
+            "env bootstrap: honoring pre-set %s=%s (overrides gin binding)",
+            name,
+            os.environ[name],
+        )
+    elif value is not None:
+        os.environ[name] = "1" if value else "0"
+        logger.info("env bootstrap: %s=%s", name, os.environ[name])
+
+
 @gin.configurable
 def apply_env_bootstrap(
     TRITON_FULL_AUTOTUNE: Optional[bool] = None,
+    TRITON_HSTU_LAST_LAYER_TARGETS_ONLY: Optional[bool] = None,
+    TRITON_HSTU_TARGETS_ONLY_EVAL: Optional[bool] = None,
 ) -> None:
     # A pre-set environment variable wins over the gin binding. The pinned
     # triton configs are MI350X-specific, so a different GPU arch (e.g. B200
@@ -36,3 +55,10 @@ def apply_env_bootstrap(
     elif TRITON_FULL_AUTOTUNE is not None:
         os.environ["TRITON_FULL_AUTOTUNE"] = "1" if TRITON_FULL_AUTOTUNE else "0"
         logger.info("env bootstrap: TRITON_FULL_AUTOTUNE=%s", os.environ["TRITON_FULL_AUTOTUNE"])
+
+    # Read at Triton/module import time, which is why they belong here rather
+    # than in a config object the model reads at construction.
+    _bind_bool(
+        "TRITON_HSTU_LAST_LAYER_TARGETS_ONLY", TRITON_HSTU_LAST_LAYER_TARGETS_ONLY
+    )
+    _bind_bool("TRITON_HSTU_TARGETS_ONLY_EVAL", TRITON_HSTU_TARGETS_ONLY_EVAL)

recommendation/generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin

@@ -32,6 +32,40 @@ make_model.hammer_kernel = "TRITON"
 # pinned constants in ops/triton/_autotune_pinning.py call sites.
 apply_env_bootstrap.TRITON_FULL_AUTOTUNE = False
 
+# In the LAST HSTU layer, compute U/Q for only the one candidate row per
+# sequence instead of all L rows. That row is the only one whose output the loss
+# consumes, so the full-length U/Q projection, the full attention, and the full
+# output projection are all discarded work; K/V stay full length because the
+# candidate attends over the entire history. Exact in math (the surviving row is
+# the same function of the same inputs, and the dropout mask is drawn with the
+# candidate's own RNG row indices so the mask stream matches the full path),
+# though not bit-identical: the compacted reductions run in a different order.
+#
+# GUARDED AT RUNTIME. It activates only for a shape where the rewrite is
+# provably equivalent — not inference, non-listwise, exactly one
+# non-interleaved suffix target per sequence, target-aware causal attention, no
+# contextual/window mask, no group norm, no full-embedding consumer, hardware
+# with separated-RNG dropout masks, and the Triton kernel path. Anything else
+# takes the normal full-length forward, so this is safe to leave on across
+# configs. See STULayer.supports_targets_only.
+#
+# NOTE FOR PERF ACCOUNTING: this removes real FLOPs. HFU prices them — it
+# divides by the `exec` budget, which carries the discount. MFU deliberately
+# keeps the full-length yardstick so it stays a fixed reference, which means it
+# RISES when this is on and must not be read as better utilization. Compare
+# ms/step or HFU, not MFU.
+# Override per-run via $TRITON_HSTU_LAST_LAYER_TARGETS_ONLY.
+apply_env_bootstrap.TRITON_HSTU_LAST_LAYER_TARGETS_ONLY = True
+
+# Whether the eval passes take that same path. The holdout batches are one
+# target per sequence just like the training ones, so the rewrite is equally
+# exact there, and eval would otherwise pay a full-length last layer for no
+# reason. Split out from the knob above only so an A/B can move eval without
+# moving training: eval numerics are what the convergence metric and the RCP
+# comparison are built on, and reduction order shifts them in the last digits.
+# Override per-run via $TRITON_HSTU_TARGETS_ONLY_EVAL.
+apply_env_bootstrap.TRITON_HSTU_TARGETS_ONLY_EVAL = True
+
 # =============================================================================
 # $SEED — global RNG seed for reproducible MODEL INITIALIZATION.
 #

recommendation/generative_recommenders/dlrm_v4/utils.py

@@ -1370,10 +1370,15 @@ def compute_and_log(
             #   tflops_algo/gpu, mfu  — uses max_seq_len^2 attention work (the
             #     MFU yardstick: the FLOPs the workload would do if every
             #     user's UIH filled the padded seq length).
-            #   tflops_real/gpu, hfu — uses this batch's mean(s_i^2) (actual
-            #     GPU work; hardware utilization).
-            #   fill                  — real / algo as a percent; how much of
-            #     the algo budget the model actually executed this batch.
+            #   tflops_real/gpu, hfu — uses this batch's mean(s_i^2) and
+            #     discounts a targets-only last layer (actual GPU work;
+            #     hardware utilization).
+            #   fill                  — ragged / algo as a percent; how much of
+            #     the padded sequence budget this batch's sequences fill.
+            #   exec                  — real / algo as a percent; how much of
+            #     the algo budget the model actually executed, so it carries
+            #     both the ragged fill and any work a knob removed outright
+            #     (exec / fill isolates the latter).
             # The jagged stash is read from the inner model; the model ref may
             # be a DMP wrapper, so unwrap via .module if present.
             tflops_str = ""
@@ -1385,20 +1390,29 @@ def compute_and_log(
                 self.tb_logger.add_scalar("perf/train_mfu_pct", mfu, global_step=step)
                 tflops_str = f" tflops_algo/gpu={tflops_algo:.1f} mfu={mfu:.1f}%"
                 jagged_t = None
+                executed_t = None
                 m = self._model_ref
                 if m is not None:
                     inner = m.module if hasattr(m, "module") else m
                     jagged_t = getattr(inner, "_last_jagged_flops_per_sample", None)
+                    executed_t = getattr(
+                        inner, "_last_executed_flops_per_sample", None
+                    )
                 if jagged_t is not None:
-                    jagged = float(jagged_t.item())
-                    if 0 < jagged < self._num_flops_per_sample:
-                        tflops_real = jagged * local_sps / 1e12
-                        hfu = 100.0 * jagged * local_sps / self._gpu_peak_flops
+                    if executed_t is None:
+                        executed_t = jagged_t
+                    # One D->H sync for both, not two.
+                    jagged, executed = torch.stack([jagged_t, executed_t]).tolist()
+                    if 0 < executed <= jagged < self._num_flops_per_sample:
+                        tflops_real = executed * local_sps / 1e12
+                        hfu = 100.0 * executed * local_sps / self._gpu_peak_flops
                         fill = 100.0 * jagged / self._num_flops_per_sample
+                        executed_pct = 100.0 * executed / self._num_flops_per_sample
                         self.tb_logger.add_scalar("perf/train_tflops_real_gpu", tflops_real, global_step=step)
                         self.tb_logger.add_scalar("perf/train_hfu_pct", hfu, global_step=step)
                         self.tb_logger.add_scalar("perf/train_fill_pct", fill, global_step=step)
-                        tflops_str += f" tflops_real/gpu={tflops_real:.1f} hfu={hfu:.1f}% fill={fill:.1f}%"
+                        self.tb_logger.add_scalar("perf/train_exec_pct", executed_pct, global_step=step)
+                        tflops_str += f" tflops_real/gpu={tflops_real:.1f} hfu={hfu:.1f}% fill={fill:.1f}% exec={executed_pct:.1f}%"
             logger.info(
                 f"train - Step {step} perf: local_sps={local_sps:.1f} "
                 f"global_sps={global_sps:.1f} step_ms={step_ms:.2f} "

recommendation/generative_recommenders/modules/dlrm_hstu.py

@@ -148,11 +148,14 @@ def __init__(  # noqa C901
         self._pipeline_mode: bool = False
         self._hstu_configs = hstu_configs
         self._bf16_training: bool = bf16_training
-        # Last batch's jagged FLOPs/sample (0-d tensor on GPU). Populated by
-        # main_forward; MetricsLogger reads + .item()s on each compute_and_log
-        # to compute tflops_real/gpu and hfu (vs dense yardstick from
-        # get_num_flops_per_sample()).
+        # Last batch's FLOPs/sample (0-d tensors on GPU), populated by
+        # main_forward and read + .item()ed by MetricsLogger on each
+        # compute_and_log, both against the dense yardstick from
+        # get_num_flops_per_sample(). `jagged` charges every layer at full
+        # length and yields `fill`; `executed` also discounts a targets-only
+        # last layer and yields tflops_real/gpu, hfu and `exec`.
         self._last_jagged_flops_per_sample: Optional[torch.Tensor] = None
+        self._last_executed_flops_per_sample: Optional[torch.Tensor] = None
         set_static_max_seq_lens([self._hstu_configs.max_seq_len])
 
         if not is_dense:
@@ -387,13 +390,44 @@ def _hstu_layer_flops(
         out = gemm * n_tokens_linear * (3 * H * hd) * D
         return uvqk + attn + out
 
+    def _hstu_layer_flops_targets_only(self, n_tokens_linear: float) -> float:
+        """Per-layer FLOPs for a layer running the targets-only path.
+
+        With ``TRITON_HSTU_LAST_LAYER_TARGETS_ONLY`` the loss consumes one
+        candidate row per sequence, so that layer computes U/Q for that row
+        alone while V/K stay full length: the ``uvqk`` gemm keeps the V and K
+        halves and drops all but one row of the U and Q halves, attention
+        becomes one query row over the whole history rather than a causal
+        triangle, and the output projection runs on that single row.
+
+        Charging this instead of ``_hstu_layer_flops`` is what keeps HFU
+        honest -- billing the layer at full length credits the model with
+        ~28% of the 3-layer total that it never executes.
+        """
+        cfg = self._hstu_configs
+        D = cfg.hstu_transducer_embedding_dim
+        H = cfg.hstu_num_heads
+        hd = cfg.hstu_attn_linear_dim
+        qd = cfg.hstu_attn_qk_dim
+        gemm = self._FLOPS_PER_MAC * self._GEMM_FWD_BWD
+        attn_mult = self._FLOPS_PER_MAC * self._ATTN_FWD_BWD
+        uvqk = gemm * (n_tokens_linear + 1.0) * D * (hd + qd) * H
+        attn = attn_mult * n_tokens_linear * H * (qd + hd)
+        out = gemm * (3 * H * hd) * D
+        return uvqk + attn + out
+
     def get_num_flops_per_sample(self) -> float:
         """Dense-equivalent fwd+bwd FLOPs per sample at ``max_seq_len``.
 
         Used as the MFU yardstick (peak utilization the workload could
         theoretically reach if every sample's sequence were the full padded
         length). The actual ``tflops_real``/``hfu`` reported per step uses
         the jagged estimate stashed by ``main_forward``.
+
+        Deliberately blind to ``TRITON_HSTU_LAST_LAYER_TARGETS_ONLY``: this is
+        the workload the model *defines*, so a knob that removes work should
+        not move it. That does mean MFU rises when the knob is on -- the
+        jagged estimate, and therefore HFU, is the one that tracks what ran.
         """
         cfg = self._hstu_configs
         S = float(cfg.max_seq_len)
@@ -416,12 +450,31 @@ def _compute_jagged_flops_per_sample(
         self,
         uih_seq_lengths: torch.Tensor,
         num_candidates: torch.Tensor,
-    ) -> torch.Tensor:
+        targets_only: bool = False,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         """Jagged fwd+bwd FLOPs per sample for THIS batch's actual lengths.
 
         Per-sample merged sequence length s_i = uih_seq_lengths[i] +
-        num_candidates[i]. Returns a 0-d tensor on the batch's device;
-        caller should ``.item()`` it (one D→H sync per logging interval).
+        num_candidates[i]. Returns ``(ragged, executed)`` as 0-d tensors on the
+        batch's device; the caller should ``.item()`` them (one D→H sync per
+        logging interval).
+
+        Two numbers because two independent things shrink the work below the
+        dense yardstick, and collapsing them would make either one unreadable:
+
+        * ``ragged`` charges every layer at full length, so against the dense
+          yardstick it isolates ragged sequences alone. This is ``fill``.
+        * ``executed`` additionally discounts the last layer when it took the
+          targets-only path, so it is the work the GPU really did. This drives
+          ``tflops_real``/``hfu``, and against the yardstick it is ``exec``.
+
+        The two are equal whenever the knob is off or the gate declined, so
+        ``exec``/``fill`` isolates the work removal on its own.
+
+        ``targets_only`` reports whether the last layer actually took the
+        targets-only path on this batch, and must come from the transducer
+        rather than the env flag: the gate declines per batch at eval, or when
+        a batch is not exactly one target per sequence.
         """
         s = (uih_seq_lengths + num_candidates).float()
         cfg = self._hstu_configs
@@ -430,18 +483,29 @@ def _compute_jagged_flops_per_sample(
         # mean(s_i²) rather than mean(s_i)² is what the attention term needs:
         # cost is quadratic per sample, so it must be averaged after squaring
         # (Jensen — using the squared mean would understate a skewed batch).
-        flops = cfg.hstu_attn_num_layers * self._hstu_layer_flops(
+        per_layer = self._hstu_layer_flops(
             n_tokens_linear=s.mean(), n_tokens_attn_sq=(s * s).mean()
         )
+        ragged = cfg.hstu_attn_num_layers * per_layer
+        if targets_only:
+            executed = (
+                cfg.hstu_attn_num_layers - 1
+            ) * per_layer + self._hstu_layer_flops_targets_only(
+                n_tokens_linear=s.mean()
+            )
+        else:
+            executed = ragged
         n_tasks = len(self._multitask_configs)
         if n_tasks > 0:
-            flops = flops + (
+            head = (
                 self._FLOPS_PER_MAC
                 * self._GEMM_FWD_BWD
                 * n_tasks
                 * cfg.hstu_transducer_embedding_dim
             )
-        return flops
+            ragged = ragged + head
+            executed = executed + head
+        return ragged, executed
 
     def _construct_payload(
         self,
@@ -701,18 +765,6 @@ def main_forward(
         Optional[torch.Tensor],
         Optional[torch.Tensor],
     ]:
-        # Stash this batch's jagged FLOPs/sample for MetricsLogger to read.
-        # No D->H sync: the .item() happens once per metric_log_frequency in
-        # the trainer, not on every step. Eval-mode batches also produce a
-        # stash but the trainer only consumes it on train batches.
-        if not torch.jit.is_scripting():
-            self._last_jagged_flops_per_sample = (
-                self._compute_jagged_flops_per_sample(
-                    uih_seq_lengths=uih_seq_lengths,
-                    num_candidates=num_candidates,
-                )
-            )
-
         # merge uih and candidates embeddings
         for (
             uih_feature_name,
@@ -751,6 +803,23 @@ def main_forward(
                 total_uih_len=total_uih_len,
                 total_targets=total_targets,
             )
+
+        # Stash this batch's jagged FLOPs/sample for MetricsLogger to read.
+        # No D->H sync: the .item() happens once per metric_log_frequency in
+        # the trainer, not on every step. Eval-mode batches also produce a
+        # stash but the trainer only consumes it on train batches. Stashed
+        # after the transducer ran so it can be priced against the path that
+        # was actually taken, not the one the env flag asked for.
+        if not torch.jit.is_scripting():
+            (
+                self._last_jagged_flops_per_sample,
+                self._last_executed_flops_per_sample,
+            ) = self._compute_jagged_flops_per_sample(
+                uih_seq_lengths=uih_seq_lengths,
+                num_candidates=num_candidates,
+                targets_only=self._hstu_transducer._last_targets_only,
+            )
+
         with record_function("## multitask_module ##"):
             supervision_labels, supervision_weights = (
                 _get_supervision_labels_and_weights(