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, + )