From e60300b8a4bf45396f2afa1a5d629cc4fe51efc8 Mon Sep 17 00:00:00 2001 From: Jiawei Chen Date: Fri, 28 Aug 2026 11:58:19 +0000 Subject: [PATCH] [DLRMv4] Shrink the last HSTU layer to the rows the loss reads The loss consumes one candidate row per sequence, so in the final HSTU layer every other row's output is discarded. This computes U and Q for that row alone against full-length K/V -- K and V stay full length because the candidate attends over the whole history -- which drops roughly 28% of three-layer FLOPs without changing what is computed. Exact in math, reordered in arithmetic: the surviving row is the same function of the same inputs, but its reductions run in a different order, so gradients are close rather than bit-identical. The dropout mask is drawn with the candidate's own RNG row indices, so the mask stream matches the full-length path rather than merely resembling it. Gated rather than assumed. Nine static conditions raise and name the blocker, since the knob is default-on and a silent fallback reads as "the optimization didn't help". Conditions that vary per batch -- a batch that is not one target per sequence, or eval when its own knob is off -- fall back and log once per distinct reason. Hardware without the separated-RNG dropout path also falls back rather than raising: below sm_100 the indexed mask does not exist, and that is not something a config change can fix. supports_indexed_output_dropout() is now the single predicate for that, used by both the gate and the kernel. The FLOPs counter needed splitting to report this honestly. It charged every layer at full length, which inflated HFU by ~28% exactly when the work was being removed. `fill` keeps its original ragged-vs-padded meaning; a new `exec` carries the discount and is what HFU divides by. MFU deliberately keeps the dense yardstick, so it RISES when this is on and must not be read as better utilization. Both knobs default on in yambda_5b.gin and are overridable by env var. --- .../dlrm_v4/train/_env_bootstrap.py | 26 ++ .../dlrm_v4/train/gin/yambda_5b.gin | 34 +++ .../generative_recommenders/dlrm_v4/utils.py | 32 +- .../modules/dlrm_hstu.py | 113 +++++-- .../modules/hstu_transducer.py | 159 +++++++++- .../generative_recommenders/modules/stu.py | 148 ++++++++- .../modules/tests/stu_targets_only_test.py | 288 ++++++++++++++++++ .../ops/hstu_compute.py | 42 +++ .../ops/triton/triton_hstu_attention.py | 212 +++++++++++++ .../ops/triton/triton_hstu_linear.py | 136 ++++++++- .../triton_hstu_preprocess_and_attention.py | 270 ++++++++++++++++ 11 files changed, 1405 insertions(+), 55 deletions(-) create mode 100644 recommendation/generative_recommenders/modules/tests/stu_targets_only_test.py diff --git a/recommendation/generative_recommenders/dlrm_v4/train/_env_bootstrap.py b/recommendation/generative_recommenders/dlrm_v4/train/_env_bootstrap.py index 5470d4e39..3f0f6ecfa 100644 --- a/recommendation/generative_recommenders/dlrm_v4/train/_env_bootstrap.py +++ b/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) diff --git a/recommendation/generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin b/recommendation/generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin index 3a1f87b19..26bca7d18 100644 --- a/recommendation/generative_recommenders/dlrm_v4/train/gin/yambda_5b.gin +++ b/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. # diff --git a/recommendation/generative_recommenders/dlrm_v4/utils.py b/recommendation/generative_recommenders/dlrm_v4/utils.py index 63a97a745..b0772bae5 100644 --- a/recommendation/generative_recommenders/dlrm_v4/utils.py +++ b/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} " diff --git a/recommendation/generative_recommenders/modules/dlrm_hstu.py b/recommendation/generative_recommenders/modules/dlrm_hstu.py index 6a667c883..635e7d139 100644 --- a/recommendation/generative_recommenders/modules/dlrm_hstu.py +++ b/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,6 +390,32 @@ 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``. @@ -394,6 +423,11 @@ def get_num_flops_per_sample(self) -> float: 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( diff --git a/recommendation/generative_recommenders/modules/hstu_transducer.py b/recommendation/generative_recommenders/modules/hstu_transducer.py index ce91a67c9..98bfefe27 100644 --- a/recommendation/generative_recommenders/modules/hstu_transducer.py +++ b/recommendation/generative_recommenders/modules/hstu_transducer.py @@ -17,7 +17,8 @@ # pyre-strict import logging -from typing import Dict, Optional, Tuple +import os +from typing import Dict, Optional, Set, Tuple import torch from generative_recommenders.common import fx_unwrap_optional_tensor, HammerModule @@ -33,6 +34,26 @@ logger: logging.Logger = logging.getLogger(__name__) torch.fx.wrap("len") +_HSTU_LAST_LAYER_TARGETS_ONLY: bool = ( + os.environ.get("TRITON_HSTU_LAST_LAYER_TARGETS_ONLY", "0") == "1" +) +# Whether eval takes the targets-only path too. Separately switchable because +# eval numerics feed the convergence metric the RCP sweep is compared against, +# so an A/B has to be able to move eval without moving training. +_HSTU_TARGETS_ONLY_EVAL: bool = ( + os.environ.get("TRITON_HSTU_TARGETS_ONLY_EVAL", "1") == "1" +) +_TARGETS_ONLY_CHECKED: bool = False +_TARGETS_ONLY_SKIPS: Set[str] = set() + + +def _log_targets_only_skip_once(reason: str) -> None: + if reason not in _TARGETS_ONLY_SKIPS: + _TARGETS_ONLY_SKIPS.add(reason) + logger.warning( + "[hstu] last layer: falling back to the full path for this batch " + f"({reason}); targets-only needs exactly one target per sequence" + ) try: torch.ops.load_library("//deeplearning/fbgemm/fbgemm_gpu:sparse_ops") @@ -78,6 +99,11 @@ def __init__( self._input_dropout_ratio: float = input_dropout_ratio self._return_full_embeddings: bool = return_full_embeddings self._listwise_training: bool = listwise and self.is_train + # Whether the last forward actually took the targets-only path. Read by + # DlrmHSTU to price the FLOPs counter, which would otherwise bill the + # last layer at full sequence length. Per-batch, because the gate can + # legitimately decline (eval, or not one target per sequence). + self._last_targets_only: bool = False for name, m in self.named_modules(): if "_stu_module" in name: @@ -168,6 +194,87 @@ def _preprocess( output_seq_payloads, ) + @torch.jit.unused + def _targets_only_unsupported_reason(self) -> Optional[str]: + """Why targets-only can never run in this configuration, or None. + + Only conditions fixed for the whole run belong here. Anything that can + differ between two batches of the same run is decided per call instead. + """ + if self._is_inference: + return "transducer is in inference mode" + if self._return_full_embeddings: + return "return_full_embeddings=True consumes every row, not just candidates" + if self._listwise_training: + return "listwise training does not pass num_targets to attention" + if self._input_preprocessor.interleave_targets(): + return "preprocessor interleaves targets, so they are not a suffix" + return self._stu_module.targets_only_unsupported_reason() + + @torch.jit.unused + def _resolve_targets_only( + self, max_targets: int, total_targets: int, batch_size: int + ) -> bool: + """Decide the last-layer targets-only path, and fail loudly if asked for + something this configuration can never do. + + The knob defaults off, so reaching here means it was explicitly turned + on. A config that cannot satisfy it is then a mistake worth raising on: + a silent fallback would read as "the optimization simply didn't help", + and nobody re-checks a knob that appears to be working. A batch that is + not exactly one target per sequence is not a mistake, so that only falls + back, logged once per distinct reason so a rare shape cannot hide. A + device without the separated-RNG dropout path falls back the same way: + that is not a mistake either, since no config change can fix it. + + Eval takes the same path as training. The two are exact in math -- with + one target per sequence the standard path's ``split_2D_jagged`` selects + precisely the row targets-only computes, and both feed it through the + same output postprocessor -- so the only difference is reduction order. + Dropout is keyed off ``self.training`` inside the layer, not off this + decision, so it stays off in eval either way. + """ + global _TARGETS_ONLY_CHECKED + if not _TARGETS_ONLY_CHECKED: + _TARGETS_ONLY_CHECKED = True + reason = self._targets_only_unsupported_reason() + if reason is not None: + raise RuntimeError( + "TRITON_HSTU_LAST_LAYER_TARGETS_ONLY=1 was requested but " + f"this configuration cannot support it: {reason}. Set " + "TRITON_HSTU_LAST_LAYER_TARGETS_ONLY=0 (or gin " + "apply_env_bootstrap.TRITON_HSTU_LAST_LAYER_TARGETS_ONLY" + "=False) to run the standard last layer instead." + ) + logger.info( + "[hstu] last layer: targets-only (one candidate row per sequence)" + ) + # Hardware, unlike the conditions above, is not something the caller can + # go fix, so this declines instead of raising. Training needs the + # candidate row to draw the same dropout mask the full-length pass would + # have drawn, and only the separated-RNG path can index a mask that way. + # Eval is unaffected: dropout is off, so no mask has to match. + from generative_recommenders.ops.triton.triton_hstu_linear import ( + supports_indexed_output_dropout, + ) + + if self.training and not supports_indexed_output_dropout(): + _log_targets_only_skip_once( + "device has no separated-RNG dropout mask path, so the compact " + "row cannot reproduce the full-length dropout stream" + ) + return False + if not self.training and not _HSTU_TARGETS_ONLY_EVAL: + _log_targets_only_skip_once("eval (TRITON_HSTU_TARGETS_ONLY_EVAL=0)") + return False + if max_targets != 1 or total_targets != batch_size: + _log_targets_only_skip_once( + f"max_targets={max_targets} total_targets={total_targets} " + f"batch_size={batch_size}" + ) + return False + return True + def _hstu_compute( self, max_seq_len: int, @@ -176,15 +283,25 @@ def _hstu_compute( seq_timestamps: torch.Tensor, seq_embeddings: torch.Tensor, num_targets: torch.Tensor, + targets_only: bool, ) -> torch.Tensor: with record_function("hstu"): - seq_embeddings = self._stu_module( - max_seq_len=max_seq_len, - x=seq_embeddings, - x_lengths=seq_lengths, - x_offsets=seq_offsets, - num_targets=(None if self._listwise_training else num_targets), - ) + if targets_only: + seq_embeddings = self._stu_module.forward_targets_only( + max_seq_len=max_seq_len, + x=seq_embeddings, + x_lengths=seq_lengths, + x_offsets=seq_offsets, + num_targets=num_targets, + ) + else: + seq_embeddings = self._stu_module( + max_seq_len=max_seq_len, + x=seq_embeddings, + x_lengths=seq_lengths, + x_offsets=seq_offsets, + num_targets=(None if self._listwise_training else num_targets), + ) return seq_embeddings def _postprocess( @@ -199,8 +316,23 @@ def _postprocess( seq_embeddings: torch.Tensor, num_targets: torch.Tensor, seq_payloads: Dict[str, torch.Tensor], + targets_only: bool, ) -> Tuple[Optional[torch.Tensor], torch.Tensor]: with record_function("hstu_output_postprocessor"): + if targets_only: + seq_offsets = torch.ops.fbgemm.asynchronous_complete_cumsum( + seq_lengths + ) + candidate_rows = seq_offsets[1:] - 1 + candidate_timestamps = torch.index_select( + seq_timestamps, 0, candidate_rows + ) + candidate_embeddings = self._output_postprocessor( + seq_embeddings=seq_embeddings, + seq_timestamps=candidate_timestamps, + seq_payloads=seq_payloads, + ) + return None, candidate_embeddings if self._return_full_embeddings: seq_embeddings = self._output_postprocessor( seq_embeddings=seq_embeddings, @@ -296,6 +428,15 @@ def forward( seq_payloads=seq_payloads, ) + targets_only = False + if _HSTU_LAST_LAYER_TARGETS_ONLY and not torch.jit.is_scripting(): + targets_only = self._resolve_targets_only( + max_targets=max_targets, + total_targets=total_targets, + batch_size=seq_lengths.size(0), + ) + if not torch.jit.is_scripting(): + self._last_targets_only = targets_only encoded_embeddings = self._hstu_compute( max_seq_len=max_seq_len, seq_lengths=seq_lengths, @@ -303,6 +444,7 @@ def forward( seq_timestamps=seq_timestamps, seq_embeddings=seq_embeddings, num_targets=num_targets, + targets_only=targets_only, ) encoded_embeddings, encoded_candidate_embeddings = self._postprocess( @@ -316,6 +458,7 @@ def forward( seq_timestamps=seq_timestamps, num_targets=num_targets, seq_payloads=seq_payloads, + targets_only=targets_only, ) if not self._is_inference: diff --git a/recommendation/generative_recommenders/modules/stu.py b/recommendation/generative_recommenders/modules/stu.py index 45c6ea5f3..89072038f 100644 --- a/recommendation/generative_recommenders/modules/stu.py +++ b/recommendation/generative_recommenders/modules/stu.py @@ -20,12 +20,17 @@ from typing import List, Optional, Tuple import torch -from generative_recommenders.common import fx_unwrap_optional_tensor, HammerModule +from generative_recommenders.common import ( + fx_unwrap_optional_tensor, + HammerKernel, + HammerModule, +) from generative_recommenders.ops.hstu_attention import delta_hstu_mha from generative_recommenders.ops.hstu_compute import ( hstu_compute_output, hstu_compute_uqvk, hstu_preprocess_and_attention, + hstu_preprocess_and_attention_targets_only, ) from generative_recommenders.ops.jagged_tensors import concat_2D_jagged, split_2D_jagged from torch.autograd.profiler import record_function @@ -39,6 +44,21 @@ class STU(HammerModule, abc.ABC): + @torch.jit.unused + def targets_only_unsupported_reason(self) -> Optional[str]: + """Why this module cannot run the targets-only path, or None if it can. + + A reason rather than a bool because the caller raises with it: the knob + is opt-in, so a config that can never satisfy it is a mistake to + surface, and "unsupported" alone does not say which of nine conditions + to go fix. + """ + return f"{type(self).__name__} has no targets-only path" + + @torch.jit.unused + def supports_targets_only(self) -> bool: + return self.targets_only_unsupported_reason() is None + def cached_forward( self, delta_x: torch.Tensor, @@ -48,6 +68,17 @@ def cached_forward( ) -> torch.Tensor: raise NotImplementedError + @torch.jit.unused + def forward_targets_only( + self, + x: torch.Tensor, + x_lengths: torch.Tensor, + x_offsets: torch.Tensor, + max_seq_len: int, + num_targets: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError + @abc.abstractmethod def forward( self, @@ -354,6 +385,82 @@ def forward( recompute_y_in_backward=self._recompute_y, ) + @torch.jit.unused + def targets_only_unsupported_reason(self) -> Optional[str]: + if self._is_inference: + return "layer is in inference mode" + if self.hammer_kernel() != HammerKernel.TRITON: + return f"hammer kernel is {self.hammer_kernel()}, not TRITON" + if self._use_group_norm: + return "group norm couples the candidate row to the rows dropped" + if not self._target_aware: + return "attention is not target-aware, so targets are not a suffix" + if not self._causal: + return "attention is not causal, so the candidate row is not last" + if self._max_attn_len != 0: + return f"max_attn_len={self._max_attn_len} windows attention" + if self._contextual_seq_len != 0: + return f"contextual_seq_len={self._contextual_seq_len} prefixes the sequence" + if not self._recompute_normed_x: + return "recompute_normed_x is off" + if not self._recompute_uvqk: + return "recompute_uvqk is off" + return None + + @torch.jit.unused + def forward_targets_only( + self, + x: torch.Tensor, + x_lengths: torch.Tensor, + x_offsets: torch.Tensor, + max_seq_len: int, + num_targets: torch.Tensor, + ) -> torch.Tensor: + """Exact final-layer training path for one suffix target per sequence.""" + if not self.supports_targets_only(): + raise RuntimeError("unsupported HSTU targets-only activation contract") + + candidate_rows = x_offsets[1:] - 1 + with record_function("## stu_targets_only_preprocess_and_attention ##"): + u, attn_output = hstu_preprocess_and_attention_targets_only( + x=x, + norm_weight=self._input_norm_weight.to(x.dtype), + norm_bias=self._input_norm_bias.to(x.dtype), + norm_eps=1e-6, + num_heads=self._num_heads, + attn_dim=self._attention_dim, + hidden_dim=self._hidden_dim, + uvqk_weight=self._uvqk_weight.to(x.dtype), + uvqk_bias=self._uvqk_beta.to(x.dtype), + max_seq_len=max_seq_len, + seq_offsets=x_offsets, + attn_alpha=self._attn_alpha, + num_targets=num_targets, + kernel=HammerKernel.TRITON, + ) + candidate_x = torch.index_select(x, 0, candidate_rows) + with record_function("## stu_targets_only_compute_output ##"): + return hstu_compute_output( + attn=attn_output, + u=u, + x=candidate_x, + norm_weight=self._output_norm_weight.to(x.dtype), + norm_bias=self._output_norm_bias.to(x.dtype), + norm_eps=1e-6, + dropout_ratio=self._output_dropout_ratio, + output_weight=self._output_weight.to(x.dtype), + group_norm=False, + num_heads=self._num_heads, + linear_dim=self._hidden_dim, + concat_u=True, + concat_x=True, + mul_u_activation_type="none", + training=self.training, + kernel=HammerKernel.TRITON, + recompute_y_in_backward=self._recompute_y, + rng_row_indices=candidate_rows, + ) + def cached_forward( self, delta_x: torch.Tensor, @@ -432,6 +539,15 @@ def __init__( super().__init__(is_inference=is_inference) self._stu_layers: torch.nn.ModuleList = torch.nn.ModuleList(modules=stu_list) + @torch.jit.unused + def targets_only_unsupported_reason(self) -> Optional[str]: + if len(self._stu_layers) == 0: + return "stack has no STU layers" + reason = self._stu_layers[ + len(self._stu_layers) - 1 + ].targets_only_unsupported_reason() # pyre-ignore [29] + return None if reason is None else f"last STU layer: {reason}" + def forward( self, x: torch.Tensor, @@ -454,6 +570,36 @@ def forward( ) return x + @torch.jit.unused + def forward_targets_only( + self, + x: torch.Tensor, + x_lengths: torch.Tensor, + x_offsets: torch.Tensor, + max_seq_len: int, + num_targets: torch.Tensor, + ) -> torch.Tensor: + num_layers = len(self._stu_layers) + if num_layers == 0: + raise RuntimeError("targets-only requires at least one STU layer") + for layer_index, layer in enumerate(self._stu_layers): + if layer_index == num_layers - 1: + return layer.forward_targets_only( # pyre-ignore [29] + x=x, + x_lengths=x_lengths, + x_offsets=x_offsets, + max_seq_len=max_seq_len, + num_targets=num_targets, + ) + x = layer( + x=x, + x_lengths=x_lengths, + x_offsets=x_offsets, + max_seq_len=max_seq_len, + num_targets=num_targets, + ) + raise AssertionError("unreachable") + def cached_forward( self, delta_x: torch.Tensor, diff --git a/recommendation/generative_recommenders/modules/tests/stu_targets_only_test.py b/recommendation/generative_recommenders/modules/tests/stu_targets_only_test.py new file mode 100644 index 000000000..761963a5e --- /dev/null +++ b/recommendation/generative_recommenders/modules/tests/stu_targets_only_test.py @@ -0,0 +1,288 @@ +#!/usr/bin/env python3 + +# pyre-strict + +import copy +import unittest +from unittest import mock + +import fbgemm_gpu # noqa: F401 +import torch +from generative_recommenders.common import ( + gpu_unavailable, + HammerKernel, + set_dev_mode, +) +from generative_recommenders.modules.stu import STULayer, STULayerConfig, STUStack + + +class STUTargetsOnlyTest(unittest.TestCase): + def _make_layer( + self, + dropout_ratio: float = 0.3, + is_inference: bool = False, + **overrides: object, + ) -> STULayer: + config = dict( + embedding_dim=64, + num_heads=2, + hidden_dim=32, + attention_dim=32, + output_dropout_ratio=dropout_ratio, + causal=True, + target_aware=True, + use_group_norm=False, + recompute_normed_x=True, + recompute_uvqk=True, + recompute_y=True, + sort_by_length=True, + contextual_seq_len=0, + ) + config.update(overrides) + layer = STULayer( + config=STULayerConfig(**config), # pyre-ignore [6] + is_inference=is_inference, + ).to(device="cuda", dtype=torch.bfloat16) + layer.recursive_setattr("_hammer_kernel", HammerKernel.TRITON) + return layer + + @unittest.skipIf(*gpu_unavailable) + def test_output_gradients_and_rng_match_full_path(self) -> None: + set_dev_mode(True) + torch.backends.cuda.matmul.allow_tf32 = False + lengths = torch.tensor([5, 9, 17, 31], device="cuda", dtype=torch.int64) + offsets = torch.ops.fbgemm.asynchronous_complete_cumsum(lengths) + num_targets = torch.ones_like(lengths) + candidate_rows = offsets[1:] - 1 + total_rows = int(offsets[-1].item()) + + reference = self._make_layer() + candidate = copy.deepcopy(reference) + torch.manual_seed(2026) + x_reference = torch.randn( + total_rows, + 64, + device="cuda", + dtype=torch.bfloat16, + ).requires_grad_() + x_candidate = x_reference.detach().clone().requires_grad_() + + torch.manual_seed(12345) + full_output = reference( + x=x_reference, + x_lengths=lengths, + x_offsets=offsets, + max_seq_len=int(lengths.max().item()), + num_targets=num_targets, + ) + expected = torch.index_select(full_output, 0, candidate_rows) + expected_post_forward_rng = torch.rand(32, device="cuda") + + torch.manual_seed(12345) + actual = candidate.forward_targets_only( + x=x_candidate, + x_lengths=lengths, + x_offsets=offsets, + max_seq_len=int(lengths.max().item()), + num_targets=num_targets, + ) + actual_post_forward_rng = torch.rand(32, device="cuda") + + torch.testing.assert_close(actual, expected, atol=5e-2, rtol=5e-2) + torch.testing.assert_close( + actual_post_forward_rng, + expected_post_forward_rng, + atol=0.0, + rtol=0.0, + ) + + torch.manual_seed(67890) + dout = torch.randn_like(expected) + expected.backward(dout) + actual.backward(dout) + torch.testing.assert_close( + x_candidate.grad, + x_reference.grad, + atol=5e-2, + rtol=5e-2, + ) + for (reference_name, reference_parameter), ( + candidate_name, + candidate_parameter, + ) in zip(reference.named_parameters(), candidate.named_parameters()): + self.assertEqual(reference_name, candidate_name) + self.assertIsNotNone(reference_parameter.grad) + self.assertIsNotNone(candidate_parameter.grad) + torch.testing.assert_close( + candidate_parameter.grad, + reference_parameter.grad, + atol=5e-2, + rtol=5e-2, + ) + + @unittest.skipIf(*gpu_unavailable) + def test_activation_contract_rejects_unsupported_layers(self) -> None: + set_dev_mode(True) + supported = self._make_layer(dropout_ratio=0.0) + self.assertTrue(supported.supports_targets_only()) + + grouped = STULayer( + config=STULayerConfig( + embedding_dim=64, + num_heads=2, + hidden_dim=32, + attention_dim=32, + use_group_norm=True, + ), + is_inference=False, + ).to(device="cuda", dtype=torch.bfloat16) + grouped.recursive_setattr("_hammer_kernel", HammerKernel.TRITON) + self.assertFalse(grouped.supports_targets_only()) + + inference = self._make_layer(dropout_ratio=0.0, is_inference=True) + self.assertFalse(inference.supports_targets_only()) + + @unittest.skipIf(*gpu_unavailable) + def test_rejection_names_the_offending_condition(self) -> None: + """A rejection has to say which condition closed the gate. + + The caller raises with this string, so "unsupported" alone would leave + someone to bisect nine conditions by hand. + """ + set_dev_mode(True) + self.assertIsNone( + self._make_layer(dropout_ratio=0.0).targets_only_unsupported_reason() + ) + + for kwargs, expected in ( + ({"causal": False}, "causal"), + ({"target_aware": False}, "target-aware"), + ({"max_attn_len": 128}, "max_attn_len"), + ({"contextual_seq_len": 4}, "contextual_seq_len"), + ({"recompute_uvqk": False}, "recompute_uvqk"), + ): + with self.subTest(**kwargs): + reason = self._make_layer( + dropout_ratio=0.0, **kwargs + ).targets_only_unsupported_reason() + self.assertIsNotNone(reason) + self.assertIn(expected, reason) + + stack = STUStack([self._make_layer(dropout_ratio=0.0, causal=False)]) + self.assertIn("causal", stack.targets_only_unsupported_reason() or "") + self.assertIn( + "no STU layers", STUStack([]).targets_only_unsupported_reason() or "" + ) + + def test_delta_backward_does_not_specialize_on_max_seq_len(self) -> None: + """max_seq_len must stay a runtime scalar in the delta backward. + + Production derives it per batch as max_uih_len + num_candidates, so as a + tl.constexpr it keys the JIT cache on a value that changes nearly every + step. A cold process then recompiles the kernel per batch -- measured at + ~100 ms each, 32 compiles for 32 distinct values -- and because the + compile blocks the launch queue it shows up as GPU idle, not host time. + A warm cache hides all of it, so no throughput or numerics test in this + suite would notice the regression. + """ + from generative_recommenders.ops.triton.triton_hstu_attention import ( + _hstu_attn_delta_bwd, + ) + + constexpr_params = { + param.name for param in _hstu_attn_delta_bwd.params if param.is_constexpr + } + self.assertNotIn("MAX_SEQ_LEN", constexpr_params) + # Only shape-independent tuning constants may be baked in: anything + # derived from the batch would reintroduce the per-step recompile. + self.assertTrue( + all(name.startswith("BLOCK_") for name in constexpr_params), + f"unexpected compile-time constant(s): {sorted(constexpr_params)}", + ) + + def test_eval_takes_the_targets_only_path(self) -> None: + """Eval must run the same shrunken last layer that training does. + + The rewrite is exact for any batch that is one target per sequence, and + the holdout batches are, so gating it on ``self.training`` only made + eval pay a full-length last layer -- worth 5.6% of eval time on an + otherwise identical pair. Guarded because the gate is one ``and`` away + from quietly excluding eval again, and no numerics test would notice. + """ + from generative_recommenders.modules import hstu_transducer as ht + + class _Stub: + training = False + + def _targets_only_unsupported_reason(self): + return None + + stub = _Stub() + with mock.patch.object(ht, "_TARGETS_ONLY_CHECKED", True): + for knob in (True, False): + with mock.patch.object(ht, "_HSTU_TARGETS_ONLY_EVAL", knob): + self.assertIs( + ht.HSTUTransducer._resolve_targets_only( + stub, max_targets=1, total_targets=4, batch_size=4 + ), + knob, + ) + # A shape the rewrite cannot express still falls back in eval. + with mock.patch.object(ht, "_HSTU_TARGETS_ONLY_EVAL", True): + self.assertIs( + ht.HSTUTransducer._resolve_targets_only( + stub, max_targets=2, total_targets=8, batch_size=4 + ), + False, + ) + + def test_device_without_indexed_dropout_declines_rather_than_raises(self) -> None: + """A device with no separated-RNG mask path must fall back, not raise. + + Training needs the candidate row to draw the same dropout mask the + full-length pass would have drawn, and only the separated-RNG path can + index a mask that way; below sm_100 (and pre-MI350) it does not exist. + Every other blocked condition here is a config the caller can go fix, so + raising is the right answer for those. Hardware is not, and raising on it + would turn a default-on knob into a hard failure on those GPUs. Eval is + exempt because dropout is off, so no mask has to match. + + This is the one gate condition that cannot be reached on the hardware + this suite runs on, so it is asserted by substitution or not at all. + """ + from generative_recommenders.modules import hstu_transducer as ht + from generative_recommenders.ops.triton import triton_hstu_linear as thl + + class _Training: + training = True + + def _targets_only_unsupported_reason(self): + return None + + class _Eval(_Training): + training = False + + with mock.patch.object(ht, "_TARGETS_ONLY_CHECKED", True), mock.patch.object( + ht, "_HSTU_TARGETS_ONLY_EVAL", True + ): + for supported, expected_in_training in ((False, False), (True, True)): + with mock.patch.object( + thl, "supports_indexed_output_dropout", lambda: supported + ): + self.assertIs( + ht.HSTUTransducer._resolve_targets_only( + _Training(), max_targets=1, total_targets=4, batch_size=4 + ), + expected_in_training, + ) + # Eval takes the path either way. + self.assertIs( + ht.HSTUTransducer._resolve_targets_only( + _Eval(), max_targets=1, total_targets=4, batch_size=4 + ), + True, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/recommendation/generative_recommenders/ops/hstu_compute.py b/recommendation/generative_recommenders/ops/hstu_compute.py index 7728c1454..a418171b2 100644 --- a/recommendation/generative_recommenders/ops/hstu_compute.py +++ b/recommendation/generative_recommenders/ops/hstu_compute.py @@ -45,6 +45,7 @@ ) from generative_recommenders.ops.triton.triton_hstu_preprocess_and_attention import ( triton_hstu_preprocess_and_attention, + triton_hstu_preprocess_and_attention_targets_only, ) from torch.fx._symbolic_trace import is_fx_tracing @@ -151,6 +152,7 @@ def hstu_compute_output( group_norm: bool, recompute_y_in_backward: bool, kernel: HammerKernel = HammerKernel.PYTORCH, + rng_row_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: if torch.jit.is_scripting(): return pytorch_hstu_compute_output( @@ -189,6 +191,7 @@ def hstu_compute_output( linear_dim=linear_dim, seed=None, recompute_y_in_backward=recompute_y_in_backward, + rng_row_indices=rng_row_indices, ) elif kernel == HammerKernel.TRITON_INFERENCE: if group_norm: @@ -388,3 +391,42 @@ def hstu_preprocess_and_attention( kernel=kernel, ).view(-1, hidden_dim * num_heads) return u, attn_output, k, v + + +def hstu_preprocess_and_attention_targets_only( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + norm_eps: float, + num_heads: int, + attn_dim: int, + hidden_dim: int, + uvqk_weight: torch.Tensor, + uvqk_bias: torch.Tensor, + max_seq_len: int, + seq_offsets: torch.Tensor, + attn_alpha: float, + num_targets: torch.Tensor, + kernel: HammerKernel = HammerKernel.TRITON, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute compact final-layer U/Q attention while retaining full K/V.""" + if not is_fx_tracing(): + torch._assert(kernel == HammerKernel.TRITON, "targets-only requires Triton") + torch._assert(x.dim() == 2, "x must be 2-D") + torch._assert(num_targets.dim() == 1, "num_targets must be 1-D") + u, attn_output = triton_hstu_preprocess_and_attention_targets_only( + x=x, + norm_weight=norm_weight, + norm_bias=norm_bias, + norm_eps=norm_eps, + num_heads=num_heads, + attn_dim=attn_dim, + hidden_dim=hidden_dim, + uvqk_weight=uvqk_weight, + uvqk_bias=uvqk_bias, + max_seq_len=max_seq_len, + seq_offsets=seq_offsets, + attn_alpha=attn_alpha, + num_targets=num_targets, + ) + return u, attn_output.view(-1, hidden_dim * num_heads) diff --git a/recommendation/generative_recommenders/ops/triton/triton_hstu_attention.py b/recommendation/generative_recommenders/ops/triton/triton_hstu_attention.py index 768ef0013..c91b70a61 100644 --- a/recommendation/generative_recommenders/ops/triton/triton_hstu_attention.py +++ b/recommendation/generative_recommenders/ops/triton/triton_hstu_attention.py @@ -2727,6 +2727,129 @@ def _hstu_attn_bwd( # noqa C901 ) +@triton.jit +def _hstu_attn_delta_bwd( + Q, + K, + V, + seq_offsets, + DOut, + DQ, + DK, + DV, + stride_qm, + stride_qh, + stride_kn, + stride_kh, + stride_vn, + stride_vh, + stride_dom, + stride_doh, + stride_dqm, + stride_dqh, + stride_dkn, + stride_dkh, + stride_dvn, + stride_dvh, + alpha, + H, + MAX_SEQ_LEN, + BLOCK_D_Q: tl.constexpr, + BLOCK_D_V: tl.constexpr, + BLOCK_N: tl.constexpr, +): + """Backward for one compact suffix query per jagged sequence. + + This specialization is intentionally narrower than the cached inference + path: every sequence has exactly one target, contextual/window modes are + disabled, and target-aware masking makes that target attend all history + rows plus itself. Each program owns one sequence/head, so dK/dV stores are + disjoint and dQ can stay in registers across the key loop. + + MAX_SEQ_LEN is a runtime scalar, not a tl.constexpr, and must stay that way. + Production derives it per batch as max_uih_len + num_candidates, so making it + constexpr keys the JIT cache on a value that changes nearly every step: a + cold process then pays a fresh ~110 ms compile per batch, and that compile + blocks the launch queue, so it shows up as GPU idle rather than as host time. + It is only ever used as the 1/MAX_SEQ_LEN scale below -- the same expression + _hstu_attn_bwd already evaluates from a runtime value -- so nothing here + benefits from constant folding. _hstu_attn_fwd/_hstu_attn_bwd handle this by + passing a separately quantized AUTOTUNE_MAX_SEQ_LEN as the autotuning key; + this kernel has a fixed config and needs no such key. + """ + off_hz = tl.program_id(0) + off_z = off_hz // H + off_h = (off_hz % H).to(tl.int64) + seq_start = tl.load(seq_offsets + off_z).to(tl.int64) + seq_end = tl.load(seq_offsets + off_z + 1) + seq_len = (seq_end - seq_start).to(tl.int32) + + offs_q = tl.arange(0, BLOCK_D_Q) + offs_v = tl.arange(0, BLOCK_D_V) + q = tl.load( + Q + off_z * stride_qm + off_h * stride_qh + offs_q, + mask=offs_q < BLOCK_D_Q, + other=0.0, + ).to(tl.float32) + dout = tl.load( + DOut + off_z * stride_dom + off_h * stride_doh + offs_v, + mask=offs_v < BLOCK_D_V, + other=0.0, + ).to(tl.float32) + dq = tl.zeros([BLOCK_D_Q], dtype=tl.float32) + + for start_n in tl.range(0, seq_len, BLOCK_N, num_stages=1): + offs_n = start_n + tl.arange(0, BLOCK_N) + mask_n = offs_n < seq_len + k_ptrs = ( + K + + (seq_start + offs_n[:, None]) * stride_kn + + off_h * stride_kh + + offs_q[None, :] + ) + v_ptrs = ( + V + + (seq_start + offs_n[:, None]) * stride_vn + + off_h * stride_vh + + offs_v[None, :] + ) + k = tl.load(k_ptrs, mask=mask_n[:, None], other=0.0) + v = tl.load(v_ptrs, mask=mask_n[:, None], other=0.0) + + qk = tl.sum(k.to(tl.float32) * q[None, :], axis=1) * alpha + sig = fast_dividef(1.0, 1.0 + tl.exp(-qk)) + silu = fast_dividef(qk, 1.0 + fast_expf(-qk)) * (1.0 / MAX_SEQ_LEN) + silu = tl.where(mask_n, silu, 0.0).to(v.dtype) + + dv = silu[:, None].to(tl.float32) * dout[None, :] + dqk = tl.sum(v.to(tl.float32) * dout[None, :], axis=1) + dqk *= sig * (1.0 + qk * (1.0 - sig)) * (1.0 / MAX_SEQ_LEN) + dqk = tl.where(mask_n, dqk, 0.0).to(k.dtype) + dk = dqk[:, None].to(tl.float32) * q[None, :] * alpha + dq += tl.sum(dqk[:, None].to(tl.float32) * k.to(tl.float32), axis=0) + + dk_ptrs = ( + DK + + (seq_start + offs_n[:, None]) * stride_dkn + + off_h * stride_dkh + + offs_q[None, :] + ) + dv_ptrs = ( + DV + + (seq_start + offs_n[:, None]) * stride_dvn + + off_h * stride_dvh + + offs_v[None, :] + ) + tl.store(dk_ptrs, dk.to(DK.dtype.element_ty), mask=mask_n[:, None]) + tl.store(dv_ptrs, dv.to(DV.dtype.element_ty), mask=mask_n[:, None]) + + tl.store( + DQ + off_z * stride_dqm + off_h * stride_dqh + offs_q, + (dq * alpha).to(DQ.dtype.element_ty), + mask=offs_q < BLOCK_D_Q, + ) + + @maybe_register_custom_op( "generative_recommenders::triton_hstu_attention_fwd", mutates_args=() ) @@ -2950,6 +3073,95 @@ def alloc_fn(size: int, align: int, stream: Optional[int]): copy_if_different_ptr(orig_dv, dv) +@maybe_register_custom_op( + "generative_recommenders::triton_hstu_attention_delta_bwd", + mutates_args=("dq", "dk", "dv"), +) +def triton_hstu_attention_delta_bwd( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + dv: torch.Tensor, + seq_offsets: torch.Tensor, + N: int, + alpha: float, +) -> None: + """Differentiate one compact suffix query per sequence against full K/V.""" + orig_dq, orig_dk, orig_dv = dq, dk, dv + dout = switch_to_contiguous_if_needed(dout) + q = switch_to_contiguous_if_needed(q) + k = switch_to_contiguous_if_needed(k) + v = switch_to_contiguous_if_needed(v) + dq = switch_to_contiguous_if_needed(dq) + dk = switch_to_contiguous_if_needed(dk) + dv = switch_to_contiguous_if_needed(dv) + if dout.shape[0] == 0: + orig_dq.zero_() + orig_dk.zero_() + orig_dv.zero_() + return + + Z = seq_offsets.numel() - 1 + _, H, DimQ = q.shape + _, _, DimV = v.shape + if q.shape[0] != Z or dout.shape[0] != Z: + raise ValueError("delta attention training requires one query per sequence") + + _hstu_attn_delta_bwd[(Z * H,)]( + Q=q, + K=k, + V=v, + seq_offsets=seq_offsets, + DOut=dout, + DQ=dq, + DK=dk, + DV=dv, + stride_qm=q.stride(0), + stride_qh=q.stride(1), + stride_kn=k.stride(0), + stride_kh=k.stride(1), + stride_vn=v.stride(0), + stride_vh=v.stride(1), + stride_dom=dout.stride(0), + stride_doh=dout.stride(1), + stride_dqm=dq.stride(0), + stride_dqh=dq.stride(1), + stride_dkn=dk.stride(0), + stride_dkh=dk.stride(1), + stride_dvn=dv.stride(0), + stride_dvh=dv.stride(1), + alpha=alpha, + H=H, + MAX_SEQ_LEN=N, + BLOCK_D_Q=DimQ, + BLOCK_D_V=DimV, + BLOCK_N=32, + num_warps=4, + ) + copy_if_different_ptr(orig_dq, dq) + copy_if_different_ptr(orig_dk, dk) + copy_if_different_ptr(orig_dv, dv) + + +@triton_hstu_attention_delta_bwd.register_fake +def _triton_hstu_attention_delta_bwd_fake( + dout: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + dq: torch.Tensor, + dk: torch.Tensor, + dv: torch.Tensor, + seq_offsets: torch.Tensor, + N: int, + alpha: float, +) -> None: + return None + + @triton_hstu_attention_fwd.register_fake def _triton_hstu_attention_fwd_fake( N: int, diff --git a/recommendation/generative_recommenders/ops/triton/triton_hstu_linear.py b/recommendation/generative_recommenders/ops/triton/triton_hstu_linear.py index 516a15664..b774d6a0b 100644 --- a/recommendation/generative_recommenders/ops/triton/triton_hstu_linear.py +++ b/recommendation/generative_recommenders/ops/triton/triton_hstu_linear.py @@ -83,6 +83,18 @@ def set_fuse_output_ln_rng_blackwell(value: bool) -> None: FUSE_OUTPUT_LN_RNG_BLACKWELL = value +def supports_indexed_output_dropout() -> bool: + """Whether ``rng_row_indices`` can be honoured on this device. + + The indexed dropout mask is only built on the separated-RNG path, so a + caller that needs a compact subset of rows to draw the same mask the + full-length pass would have drawn must check this first and take the + full-length path when it is False. Passing indices anyway is rejected + below rather than silently drawing a different mask. + """ + return not FUSE_OUTPUT_LN_RNG_BLACKWELL and use_separated_rng_ln_mul_dropout() + + @triton.jit def rand3x(seed, offsets, n_rounds: tl.constexpr = 10): # pyre-ignore [9] i1, i2, i3, _ = tl.randint4x(seed, offsets, n_rounds) @@ -141,6 +153,43 @@ def _generate_random_mask( tl.store(base_ptr + 3 * STRIDE, packed3, mask=row3_mask) +@triton.jit +def _generate_indexed_random_mask( + MASK_BUFFER, + ROW_INDICES, + N, + dropout_ratio, + seed, + D: tl.constexpr, + STRIDE: tl.constexpr, + BLOCK_D: tl.constexpr, + NUM_MASKS: tl.constexpr, +): + """Generate masks for compact rows using their full-tensor RNG row IDs.""" + row = tl.program_id(0) + cols = tl.arange(0, BLOCK_D) + col_mask = cols < D + logical_row = tl.load(ROW_INDICES + row).to(tl.int64) + philox_group = logical_row // 4 + philox_lane = logical_row % 4 + rand_offset = philox_group * (NUM_MASKS * BLOCK_D) + cols + packed = tl.zeros([BLOCK_D], dtype=tl.int8) + for j in tl.static_range(NUM_MASKS): + r0, r1, r2, r3 = tl.rand4x(seed, rand_offset) + selected = tl.where( + philox_lane == 0, + r0, + tl.where(philox_lane == 1, r1, tl.where(philox_lane == 2, r2, r3)), + ) + packed |= (selected > dropout_ratio).to(tl.int8) << j + rand_offset += BLOCK_D + tl.store( + MASK_BUFFER + row.to(tl.int64) * STRIDE + cols, + packed, + mask=(row < N) & col_mask, + ) + + @triton_autotune( configs=_get_layer_norm_mul_dropout_fwd_multirow_configs(), key=["BLOCK_D"], @@ -1021,6 +1070,34 @@ def _create_dropout_mask( return random_mask +def _create_indexed_dropout_mask( + row_indices: torch.Tensor, + N: int, + D: int, + BLOCK_D: int, + concat_u: bool, + concat_x: bool, + dropout_ratio: float, + seed: int, + device: torch.device, +) -> torch.Tensor: + """Create compact masks identical to selected rows of a full mask.""" + num_masks = 1 + int(concat_u) + int(concat_x) + random_mask = torch.empty([N, D], dtype=torch.int8, device=device) + _generate_indexed_random_mask[(N,)]( + random_mask, + row_indices, + N, + dropout_ratio, + seed, + D, # pyre-ignore[6] + random_mask.stride(0), # pyre-ignore[6] + BLOCK_D, # pyre-fixme[6]: Triton constexpr param + num_masks, # pyre-ignore[6]: NUM_MASKS constexpr + ) + return random_mask + + @maybe_register_custom_op( "generative_recommenders::_triton_layer_norm_mul_dropout_fwd_impl", mutates_args=() ) @@ -1037,6 +1114,7 @@ def _triton_layer_norm_mul_dropout_fwd_impl( concat_x: bool, mul_u_activation_type: str, seed: int, + rng_row_indices: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Internal implementation that returns only tensors for custom_op compatibility. @@ -1069,21 +1147,30 @@ def _triton_layer_norm_mul_dropout_fwd_impl( # backward, instead of launching one program per row with fused RNG. This is a large # win on Blackwell (sm_100) and AMD MI350 (gfx950); other GPUs keep the fused path. # Extended to support concat_u + concat_x for mask reuse optimization - if ( - not FUSE_OUTPUT_LN_RNG_BLACKWELL - and use_separated_rng_ln_mul_dropout() - and training - ): - random_mask = _create_dropout_mask( - N=N, - D=D, - BLOCK_D=BLOCK_D, - concat_u=concat_u, - concat_x=concat_x, - dropout_ratio=dropout_ratio, - seed=seed, - device=x.device, - ) + if supports_indexed_output_dropout() and training: + if rng_row_indices.numel() > 0: + random_mask = _create_indexed_dropout_mask( + row_indices=rng_row_indices, + N=N, + D=D, + BLOCK_D=BLOCK_D, + concat_u=concat_u, + concat_x=concat_x, + dropout_ratio=dropout_ratio, + seed=seed, + device=x.device, + ) + else: + random_mask = _create_dropout_mask( + N=N, + D=D, + BLOCK_D=BLOCK_D, + concat_u=concat_u, + concat_x=concat_x, + dropout_ratio=dropout_ratio, + seed=seed, + device=x.device, + ) def grid(META): return (triton.cdiv(N, META["BLOCK_N"]),) @@ -1115,6 +1202,8 @@ def grid(META): ) else: + if rng_row_indices.numel() > 0 and training: + raise RuntimeError("indexed output dropout requires separated RNG masks") # Default path: fused RNG generation # Mask cannot be saved with fused RNG - it's generated inline in the kernel # pyre-ignore[28] @@ -1159,6 +1248,7 @@ def _triton_layer_norm_mul_dropout_fwd_impl_fake( concat_x: bool, mul_u_activation_type: str, seed: int, + rng_row_indices: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: """Fake implementation for FakeTensor tracing.""" N, D = x.shape @@ -1189,6 +1279,7 @@ def triton_layer_norm_mul_dropout_fwd( concat_x: bool = False, mul_u_activation_type: str = "none", seed: Optional[int] = None, + rng_row_indices: Optional[torch.Tensor] = None, ) -> Tuple[ torch.Tensor, torch.Tensor, torch.Tensor, int, int, int, Optional[torch.Tensor] ]: # y, mean, rstd, BLOCK_D, num_warps, seed, random_mask @@ -1221,6 +1312,9 @@ def triton_layer_norm_mul_dropout_fwd( assert bias.dim() == 1 assert weight.numel() == D assert bias.numel() == D + if rng_row_indices is not None: + assert rng_row_indices.dim() == 1 + assert rng_row_indices.numel() == N if N == 0: D = x.shape[1] @@ -1252,6 +1346,11 @@ def triton_layer_norm_mul_dropout_fwd( num_warps: int = min(max(BLOCK_D // 256, 1), 8) # Call internal implementation + rng_row_indices_tensor = ( + rng_row_indices + if rng_row_indices is not None + else torch.empty(0, dtype=torch.int64, device=x.device) + ) y, mean, rstd, random_mask_tensor = _triton_layer_norm_mul_dropout_fwd_impl( x, u, @@ -1265,6 +1364,7 @@ def triton_layer_norm_mul_dropout_fwd( concat_x, mul_u_activation_type, seed if seed is not None else 0, + rng_row_indices_tensor, ) # Convert empty tensor back to None @@ -2231,6 +2331,7 @@ def forward( linear_dim: int = -1, seed: Optional[int] = None, recompute_y_in_backward: bool = False, + rng_row_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: if dropout_ratio == 0.0: training = False @@ -2268,6 +2369,7 @@ def forward( concat_u=concat_u, concat_x=concat_x, seed=seed, + rng_row_indices=rng_row_indices, ) ) @@ -2323,6 +2425,7 @@ def backward( None, # linear_dim None, # seed None, # recompute_y_in_backward + None, # rng_row_indices ]: attn, u, norm_weight, norm_bias, mean, rstd, output_weight = ctx.saved_tensors[ :7 @@ -2410,6 +2513,7 @@ def backward( None, # linear_dim None, # seed None, # recompute_y_in_backward + None, # rng_row_indices ) @@ -3024,6 +3128,7 @@ def triton_hstu_compute_output( linear_dim: int = -1, seed: Optional[int] = None, recompute_y_in_backward: bool = False, + rng_row_indices: Optional[torch.Tensor] = None, ) -> torch.Tensor: return HSTUComputeOutputFunction.apply( attn, @@ -3044,4 +3149,5 @@ def triton_hstu_compute_output( linear_dim, seed, recompute_y_in_backward, + rng_row_indices, ) diff --git a/recommendation/generative_recommenders/ops/triton/triton_hstu_preprocess_and_attention.py b/recommendation/generative_recommenders/ops/triton/triton_hstu_preprocess_and_attention.py index bda97ff96..bdbe2e653 100644 --- a/recommendation/generative_recommenders/ops/triton/triton_hstu_preprocess_and_attention.py +++ b/recommendation/generative_recommenders/ops/triton/triton_hstu_preprocess_and_attention.py @@ -26,7 +26,9 @@ ) from generative_recommenders.ops.triton.triton_hstu_attention import ( _should_enable_tma, + triton_cached_hstu_mha, triton_hstu_attention_bwd, + triton_hstu_attention_delta_bwd, triton_hstu_attention_fwd, ) from generative_recommenders.ops.triton.triton_layer_norm import ( @@ -292,6 +294,238 @@ def backward( ) +class _HSTUPreprocessAndAttentionTargetsOnlyFunction(torch.autograd.Function): + """Final-layer training path with full K/V and one compact U/Q per sequence.""" + + @staticmethod + # pyre-ignore [14] + def forward( + ctx, # pyre-ignore [2] + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + norm_eps: float, + num_heads: int, + attn_dim: int, + hidden_dim: int, + uvqk_weight: torch.Tensor, + uvqk_bias: torch.Tensor, + max_seq_len: int, + seq_offsets: torch.Tensor, + attn_alpha: float, + num_targets: torch.Tensor, + enable_tma: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + normed_x, x_mean, x_rstd = triton_weighted_layer_norm_fwd( + x=x, + weight=norm_weight, + bias=norm_bias, + eps=norm_eps, + ) + BLOCK_D = compute_BLOCK_D(x) + u_dim = hidden_dim * num_heads + v_dim = hidden_dim * num_heads + q_dim = attn_dim * num_heads + k_dim = attn_dim * num_heads + w_u, w_v, w_q, w_k = uvqk_weight.split( + [u_dim, v_dim, q_dim, k_dim], dim=1 + ) + b_u, b_v, b_q, b_k = uvqk_bias.split( + [u_dim, v_dim, q_dim, k_dim], dim=0 + ) + candidate_rows = seq_offsets[1:] - 1 + + # Packing only the small weights keeps the large activation projection + # at [L, 2D] without changing checkpoint layout. + w_vk = torch.cat((w_v, w_k), dim=1) + b_vk = torch.cat((b_v, b_k), dim=0) + vk = maybe_triton_addmm_fwd(x=normed_x, w=w_vk, y=b_vk).contiguous() + v, k = vk.split([v_dim, k_dim], dim=1) + candidate_normed_x = torch.index_select(normed_x, 0, candidate_rows) + w_uq = torch.cat((w_u, w_q), dim=1) + b_uq = torch.cat((b_u, b_q), dim=0) + uq = maybe_triton_addmm_fwd( + x=candidate_normed_x, w=w_uq, y=b_uq + ).contiguous() + u, q = uq.split([u_dim, q_dim], dim=1) + q = q.view(-1, num_heads, attn_dim) + k = k.view(-1, num_heads, attn_dim) + v = v.view(-1, num_heads, hidden_dim) + out = triton_cached_hstu_mha( + N=max_seq_len, + alpha=attn_alpha, + delta_q=q, + k=k, + v=v, + seq_offsets=seq_offsets, + num_targets=num_targets, + max_attn_len=0, + contextual_seq_len=0, + enable_tma=enable_tma, + ) + + ctx.save_for_backward( + x, + norm_weight, + norm_bias, + x_mean, + x_rstd, + uvqk_weight, + uvqk_bias, + seq_offsets, + ) + ctx.attn_alpha = attn_alpha + ctx.max_seq_len = max_seq_len + ctx.hidden_dim = hidden_dim + ctx.attn_dim = attn_dim + ctx.num_heads = num_heads + ctx.norm_eps = norm_eps + ctx.norm_BLOCK_D = BLOCK_D + return F.silu(u), out + + @staticmethod + # pyre-ignore [14] + def backward( + ctx, # pyre-ignore [2] + dsilu_u: torch.Tensor, + dout: torch.Tensor, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + None, + None, + None, + None, + torch.Tensor, + torch.Tensor, + None, + None, + None, + None, + None, + ]: + ( + x, + norm_weight, + norm_bias, + x_mean, + x_rstd, + uvqk_weight, + uvqk_bias, + seq_offsets, + ) = ctx.saved_tensors + normed_x, _, _ = triton_weighted_layer_norm_fwd( + x=x, + weight=norm_weight, + bias=norm_bias, + eps=ctx.norm_eps, + mean=x_mean, + rstd=x_rstd, + ) + u_dim = ctx.hidden_dim * ctx.num_heads + v_dim = ctx.hidden_dim * ctx.num_heads + q_dim = ctx.attn_dim * ctx.num_heads + k_dim = ctx.attn_dim * ctx.num_heads + w_u, w_v, w_q, w_k = uvqk_weight.split( + [u_dim, v_dim, q_dim, k_dim], dim=1 + ) + b_u, b_v, b_q, b_k = uvqk_bias.split( + [u_dim, v_dim, q_dim, k_dim], dim=0 + ) + candidate_rows = seq_offsets[1:] - 1 + candidate_normed_x = torch.index_select(normed_x, 0, candidate_rows) + w_vk = torch.cat((w_v, w_k), dim=1) + b_vk = torch.cat((b_v, b_k), dim=0) + w_uq = torch.cat((w_u, w_q), dim=1) + b_uq = torch.cat((b_u, b_q), dim=0) + vk = maybe_triton_addmm_fwd(x=normed_x, w=w_vk, y=b_vk).contiguous() + v, k = vk.split([v_dim, k_dim], dim=1) + uq = maybe_triton_addmm_fwd( + x=candidate_normed_x, w=w_uq, y=b_uq + ).contiguous() + u, q = uq.split([u_dim, q_dim], dim=1) + + q = q.view(-1, ctx.num_heads, ctx.attn_dim) + k = k.view(-1, ctx.num_heads, ctx.attn_dim) + v = v.view(-1, ctx.num_heads, ctx.hidden_dim) + dq = torch.empty_like(q) + # Write dV/dK directly in the packed column order consumed by the + # full-row projection backward. The delta kernel honors row strides, + # so this avoids an O(LD) torch.cat materialization. + dvk = torch.empty( + (normed_x.shape[0], v_dim + k_dim), + dtype=normed_x.dtype, + device=normed_x.device, + ) + dv = dvk[:, :v_dim].view(-1, ctx.num_heads, ctx.hidden_dim) + dk = dvk[:, v_dim:].view(-1, ctx.num_heads, ctx.attn_dim) + triton_hstu_attention_delta_bwd( + dout=dout, + q=q, + k=k, + v=v, + dq=dq, + dk=dk, + dv=dv, + seq_offsets=seq_offsets, + N=ctx.max_seq_len, + alpha=ctx.attn_alpha, + ) + + du = torch.empty_like(u) + torch.ops.aten.silu_backward(dsilu_u, u, grad_input=du) + duq = torch.cat((du, dq.flatten(1, 2)), dim=1) + d_candidate_x, d_w_uq, d_b_uq = triton_addmm_bwd( + x=candidate_normed_x, + w=w_uq, + dz=duq, + is_y_1d=True, + ) + d_normed_x, d_w_vk, d_b_vk = triton_addmm_bwd( + x=normed_x, + w=w_vk, + dz=dvk, + is_y_1d=True, + ) + d_normed_x.index_add_(0, candidate_rows, d_candidate_x) + + d_w_u, d_w_q = d_w_uq.split([u_dim, q_dim], dim=1) + d_w_v, d_w_k = d_w_vk.split([v_dim, k_dim], dim=1) + d_uvqk_weight = torch.cat((d_w_u, d_w_v, d_w_q, d_w_k), dim=1) + d_b_u, d_b_q = d_b_uq.split([u_dim, q_dim], dim=0) + d_b_v, d_b_k = d_b_vk.split([v_dim, k_dim], dim=0) + d_uvqk_bias = torch.cat((d_b_u, d_b_v, d_b_q, d_b_k), dim=0) + + d_x, d_norm_weight, d_norm_bias = triton_weighted_layer_norm_bwd( + dy=d_normed_x, + x=x, + weight=norm_weight, + bias=norm_bias, + mean=x_mean, + rstd=x_rstd, + learnable=True, + eps=ctx.norm_eps, + BLOCK_D=ctx.norm_BLOCK_D, + ) + return ( + d_x, + d_norm_weight, + d_norm_bias, + None, + None, + None, + None, + d_uvqk_weight, + d_uvqk_bias, + None, + None, + None, + None, + None, + ) + + def triton_hstu_preprocess_and_attention( x: torch.Tensor, norm_weight: torch.Tensor, @@ -340,3 +574,39 @@ def triton_hstu_preprocess_and_attention( sort_by_length, enable_tma, ) + + +def triton_hstu_preprocess_and_attention_targets_only( + x: torch.Tensor, + norm_weight: torch.Tensor, + norm_bias: torch.Tensor, + norm_eps: float, + num_heads: int, + attn_dim: int, + hidden_dim: int, + uvqk_weight: torch.Tensor, + uvqk_bias: torch.Tensor, + max_seq_len: int, + seq_offsets: torch.Tensor, + attn_alpha: float, + num_targets: torch.Tensor, + enable_tma: Optional[bool] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + if enable_tma is None: + enable_tma = _should_enable_tma() + return _HSTUPreprocessAndAttentionTargetsOnlyFunction.apply( + x, + norm_weight, + norm_bias, + norm_eps, + num_heads, + attn_dim, + hidden_dim, + uvqk_weight, + uvqk_bias, + max_seq_len, + seq_offsets, + attn_alpha, + num_targets, + enable_tma, + )