[feat] ULTRA-HSTU: FP8 attention via fbgemm_gpu_hstu wheel

docs/source/models/dlrm_hstu.md

@@ -168,7 +168,7 @@ model_config {
 - kernel: 算子实现，可选TRITON/PYTORCH/CUTLASS
 
   - TRITON: 基于Triton的实现，通常比PYTORCH快2-3x，节省2-3x显存
-  - CUTLASS: 基于CUTLASS的CUDA融合算子实现，需安装fbgemm_gpu_hstu包（DEVICE可选cu126/cu129，对应DEVICE_DOTTED为cu12.6/cu12.9：`pip install fbgemm_gpu_hstu==0.1.0+${DEVICE_DOTTED} -f https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/${DEVICE}/repo.html`），要求`attention_dim`等于`hidden_dim`，支持Ampere/Ada/Hopper GPU
+  - CUTLASS: 基于CUTLASS的CUDA融合算子实现，需安装fbgemm_gpu_hstu包（DEVICE可选cu126/cu129：`pip install fbgemm_gpu_hstu==0.1.0+20260528.5f13f139.${DEVICE} -f https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/${DEVICE}/repo.html`），要求`attention_dim`等于`hidden_dim`，支持Ampere/Ada/Hopper/Blackwell GPU
   - PYTORCH: 纯PyTorch实现，兼容性最好
 
 ### MTGR Style 配置方式

requirements/extra.txt

@@ -1,7 +1,7 @@
 dynamicemb @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/cu129/dynamicemb-0.1.0%2B20260519.e0c1fbb.cu129-cp310-cp310-linux_x86_64.whl ; python_version=="3.10"
 dynamicemb @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/cu129/dynamicemb-0.1.0%2B20260519.e0c1fbb.cu129-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
 dynamicemb @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/dynamicemb/cu129/dynamicemb-0.1.0%2B20260519.e0c1fbb.cu129-cp312-cp312-linux_x86_64.whl ; python_version=="3.12"
-fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2Bcu12.9-cp310-cp310-linux_x86_64.whl ; python_version=="3.10"
-fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2Bcu12.9-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
-fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2Bcu12.9-cp312-cp312-linux_x86_64.whl ; python_version=="3.12"
+fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2B20260528.5f13f139.cu129-cp310-cp310-linux_x86_64.whl ; python_version=="3.10"
+fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2B20260528.5f13f139.cu129-cp311-cp311-linux_x86_64.whl ; python_version=="3.11"
+fbgemm_gpu_hstu @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/hstu/cu129/fbgemm_gpu_hstu-0.1.0%2B20260528.5f13f139.cu129-cp312-cp312-linux_x86_64.whl ; python_version=="3.12"
 torch_fx_tool @ https://tzrec.oss-accelerate.aliyuncs.com/third_party/rtp/torch_fx_tool-0.0.1%2B20251201.8c109c4-py3-none-any.whl

tzrec/modules/gr/stu.py

@@ -250,6 +250,11 @@ class STULayer(STU):
             ``max_seq_len`` (legacy behavior). When set to a fixed value
             (typically the model's config ``max_seq_len``), attention output
             is invariant to batch-level seq-length.
+        fp8_quant_mode (int): FP8 attention quant mode (``-1`` = off,
+            ``0..5`` select an FP8 mode). Only honored on the training/eval
+            forward with ``Kernel.CUTLASS`` on SM90 (Hopper) or SM120
+            (Blackwell RTX, ``quant_mode=2`` only); the cached/delta
+            serving path always runs in bf16/fp16.
         is_inference (bool): whether to run in inference mode.
     """
 
@@ -278,6 +283,7 @@ def __init__(
         sla_k1: int = 0,
         sla_k2: int = 0,
         scaling_seqlen: int = -1,
+        fp8_quant_mode: int = -1,
         is_inference: bool = False,
     ) -> None:
         super().__init__(
@@ -302,6 +308,7 @@ def __init__(
         self._sla_k1: int = sla_k1
         self._sla_k2: int = sla_k2
         self._scaling_seqlen: int = scaling_seqlen
+        self._fp8_quant_mode: int = fp8_quant_mode
 
         self._uvqk_weight: torch.nn.Parameter = torch.nn.Parameter(
             torch.empty(
@@ -530,6 +537,7 @@ def forward(
                 enable_tma=self._enable_tma,
                 attn_func=local_attn_func,
                 scaling_seqlen=self._scaling_seqlen,
+                fp8_quant_mode=self._fp8_quant_mode,
             )
 
         self.update_kv_cache(
@@ -578,6 +586,11 @@ def cached_forward(
 
         Returns:
             torch.Tensor: output sequence embedding tensor.
+
+        Note:
+            This path runs in bf16/fp16 regardless of ``fp8_quant_mode``;
+            FP8 attention applies only to the full training/eval ``forward``
+            (``delta_hstu_mha`` has no FP8 kernel).
         """
         with record_function("## stu_compute_uqvk ##"):
             delta_u, delta_q, delta_k, delta_v = hstu_compute_uqvk(

tzrec/ops/_cuda/cutlass_hstu_attention.py

@@ -12,6 +12,7 @@
 from typing import Optional
 
 import torch
+from torch.fx._symbolic_trace import is_fx_tracing
 
 from tzrec.utils.logging_util import logger
 
@@ -26,6 +27,117 @@
     hstu_attn_varlen_func = None  # type: ignore[assignment]
 
 _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
+# Highest fp8_quant_mode the wheel accepts (0..5 are FP8 modes; -1 = off).
+_MAX_FP8_QUANT_MODE = 5
+
+
+def _assert_fp8_capable(fp8_quant_mode: int) -> None:
+    """Raise unless the (arch, quant_mode) pair has an FP8 kernel in the wheel.
+
+    The wheel's dispatcher (``cuda_hstu_attention.HstuAttnVarlenFunc.forward``)
+    only routes to an FP8 kernel for SM90 (Hopper, modes 0..5) and SM120
+    (Blackwell RTX, mode 2 only -- forward-only). SM80 (Ampere) and SM100
+    (Blackwell datacenter) have no FP8 kernel; on SM120 with quant_mode != 2
+    the dispatcher silently falls through to the SM80 bf16/fp16 path, so we
+    must reject that here.
+    """
+    if not torch.cuda.is_available():
+        raise RuntimeError(
+            "FP8 HSTU attention (fp8_quant_mode >= 0) requires a CUDA SM90 "
+            "(Hopper) or SM120 (Blackwell RTX) GPU; no CUDA device is available."
+        )
+    major = torch.cuda.get_device_capability()[0]
+    if major == 9:
+        return  # Hopper: all modes 0..5 supported (fwd + bwd).
+    if major == 12:
+        if fp8_quant_mode != 2:
+            raise RuntimeError(
+                "FP8 HSTU attention on SM120 (Blackwell RTX) only supports "
+                f"fp8_quant_mode=2 (per-block, forward-only); got "
+                f"fp8_quant_mode={fp8_quant_mode}."
+            )
+        return
+    raise RuntimeError(
+        "FP8 HSTU attention requires SM90 (Hopper) or SM120 (Blackwell RTX); "
+        f"got device capability major version {major}."
+    )
+
+
+@torch.library.custom_op("tzrec::cutlass_hstu_fp8_fwd", mutates_args=())
+def cutlass_hstu_fp8_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    num_contexts: Optional[torch.Tensor],
+    num_targets: Optional[torch.Tensor],
+    max_seq_len: int,
+    scaling_seqlen: int,
+    window_size_left: int,
+    window_size_right: int,
+    alpha: float,
+    fp8_quant_mode: int,
+    attn_func: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """Opaque FP8 HSTU forward for torch.export / AOTInductor.
+
+    The wheel's ``hstu_attn_varlen_func`` quantizes q/k/v in Python before
+    dispatching to the registered ``fbgemm::hstu_varlen_fwd_90`` op. For
+    block/head/batch/tensor FP8 modes that quantization iterates over
+    per-sample lengths derived from ``cu_seqlens`` and produces descale
+    tensors whose block-count dim is data-dependent. Tracing through it
+    raises ``GuardOnDataDependentSymNode`` (export) or, once the quantizer
+    is itself wrapped, leaves unbacked-SymInt-shaped descales in the graph
+    that AOTI mis-sizes at runtime. Wrapping the whole forward as one custom
+    op keeps all quantization + descales internal: the exported graph sees
+    only this op returning a tensor of ``v``'s shape, and the real Python
+    quantization runs unchanged inside the op at predict time.
+
+    Only used on the export path (see ``cutlass_hstu_mha``); eager training
+    keeps calling ``hstu_attn_varlen_func`` directly so autograd is intact.
+    """
+    return hstu_attn_varlen_func(
+        q,
+        k,
+        v,
+        cu_seqlens,
+        cu_seqlens,
+        seqused_q=None,
+        seqused_k=None,
+        max_seqlen_q=max_seq_len,
+        max_seqlen_k=max_seq_len,
+        scaling_seqlen=scaling_seqlen,
+        num_contexts=num_contexts,
+        num_targets=num_targets,
+        target_group_size=1,
+        window_size=(window_size_left, window_size_right),
+        alpha=alpha,
+        rab=None,
+        has_drab=False,
+        func=attn_func,
+        quant_mode=fp8_quant_mode,
+    )
+
+
+@cutlass_hstu_fp8_fwd.register_fake
+def _cutlass_hstu_fp8_fwd_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    num_contexts: Optional[torch.Tensor],
+    num_targets: Optional[torch.Tensor],
+    max_seq_len: int,
+    scaling_seqlen: int,
+    window_size_left: int,
+    window_size_right: int,
+    alpha: float,
+    fp8_quant_mode: int,
+    attn_func: Optional[torch.Tensor],
+) -> torch.Tensor:
+    # FP8 attention returns a bf16/fp16 output of v's shape (the kernel
+    # dequantizes internally). No data-dependent dims escape this op.
+    return v.new_empty(v.shape, dtype=v.dtype)
 
 
 @torch.fx.wrap
@@ -42,6 +154,7 @@ def cutlass_hstu_mha(
     contextual_seq_len: int = 0,
     attn_func: Optional[torch.Tensor] = None,
     scaling_seqlen: int = -1,
+    fp8_quant_mode: int = -1,
 ) -> torch.Tensor:
     """CUTLASS-based HSTU multi-head attention.
 
@@ -82,10 +195,23 @@ def cutlass_hstu_mha(
         scaling_seqlen: divisor used to scale the attention output inside
             the kernel. ``-1`` (default) falls back to ``max_seq_len`` so
             the behavior matches the legacy code path.
+        fp8_quant_mode: FP8 quantization mode forwarded to the wheel.
+            ``-1`` (default) keeps q/k/v in bf16/fp16 (no FP8). ``0..5``
+            select an FP8 mode (0 per-tensor, 1 two-direction, 2 per-block,
+            3 per-head, 4 per-batch, 5 global); the wheel quantizes q/k/v
+            internally. FP8 requires SM90 (Hopper, all modes, fwd+bwd) or
+            SM120 (Blackwell RTX, ``quant_mode=2`` only, forward-only).
 
     Returns:
         output tensor of shape (total, nheads, hidden_dim).
     """
+    if fp8_quant_mode < -1 or fp8_quant_mode > _MAX_FP8_QUANT_MODE:
+        raise ValueError(
+            f"fp8_quant_mode must be in [-1, {_MAX_FP8_QUANT_MODE}]; "
+            f"got {fp8_quant_mode}."
+        )
+    if fp8_quant_mode >= 0 and not is_fx_tracing():
+        _assert_fp8_capable(fp8_quant_mode)
     if hstu_attn_varlen_func is None:
         raise RuntimeError(
             "fbgemm_gpu_hstu wheel is not installed; cannot run CUTLASS "
@@ -167,6 +293,27 @@ def cutlass_hstu_mha(
     if scaling_seqlen == -1:
         scaling_seqlen = max_seq_len
 
+    if fp8_quant_mode >= 0 and torch.compiler.is_exporting():
+        # Route FP8 through the opaque custom op so torch.export / AOTI never
+        # trace the wheel's Python quantizer (whose data-dependent descales
+        # otherwise break export tracing / AOTI runtime). Eager train/eval
+        # falls through to hstu_attn_varlen_func below (autograd intact).
+        return torch.ops.tzrec.cutlass_hstu_fp8_fwd(
+            q,
+            k,
+            v,
+            cu_seqlens,
+            num_contexts_tensor,
+            num_targets_int32,
+            max_seq_len,
+            scaling_seqlen,
+            window_size_left,
+            window_size_right,
+            alpha,
+            fp8_quant_mode,
+            attn_func,
+        )
+
     return hstu_attn_varlen_func(
         q,
         k,
@@ -186,4 +333,5 @@ def cutlass_hstu_mha(
         rab=None,
         has_drab=False,
         func=attn_func,
+        quant_mode=fp8_quant_mode,
     )

tzrec/ops/hstu_attention.py

@@ -59,6 +59,7 @@ def hstu_mha(
     enable_tma: bool = False,
     attn_func: Optional[torch.Tensor] = None,
     scaling_seqlen: int = -1,
+    fp8_quant_mode: int = -1,
 ) -> torch.Tensor:
     """HSTU multi-head attention with kernel backend dispatch.
 
@@ -86,6 +87,10 @@ def hstu_mha(
         scaling_seqlen: divisor used to scale the attention output inside
             the kernel. ``-1`` (default) falls back to ``max_seq_len`` so
             the behavior matches the legacy code path.
+        fp8_quant_mode: FP8 quantization mode (``-1`` = off, ``0..5`` select
+            an FP8 mode). Only supported on ``Kernel.CUTLASS`` running on
+            SM90 (Hopper) or SM120 (Blackwell RTX, ``quant_mode=2`` only);
+            a value ``>= 0`` with any other kernel raises.
 
     Returns:
         output tensor of shape (total, nheads, hidden_dim).
@@ -107,6 +112,13 @@ def hstu_mha(
             "Kernel.PYTORCH for the reference implementation."
         )
 
+    if fp8_quant_mode >= 0 and kernel != Kernel.CUTLASS:
+        raise ValueError(
+            f"fp8_quant_mode={fp8_quant_mode} (FP8) is only supported on "
+            f"Kernel.CUTLASS; got kernel={kernel}. FP8 attention requires "
+            "the CUTLASS SM90 kernel."
+        )
+
     if kernel == Kernel.CUTLASS and attn_func is None:
         # Without an arbitrary mask, the CUTLASS kernel's local-window path
         # (Is_local) cannot combine with context/target masking — fall back
@@ -141,6 +153,7 @@ def hstu_mha(
             contextual_seq_len=contextual_seq_len,
             attn_func=attn_func,
             scaling_seqlen=scaling_seqlen,
+            fp8_quant_mode=fp8_quant_mode,
         )
 
     if kernel == Kernel.TRITON:

tzrec/ops/hstu_attention_test.py

@@ -100,6 +100,7 @@ def test_attn(
     rtol: Optional[float] = None,
     enable_tma: bool = False,
     scaling_seqlen: int = -1,
+    fp8_quant_mode: int = -1,
 ) -> None:
     # has_max_attn_len=True and enable_tma=True will result in TritonGPUCoalesce error
     # include/llvm/llvm/ADT/SmallVector.h:296: const_reference llvm::SmallVectorTemplateCommon<long>::operator[](size_type) const [T = long]: Assertion `idx < size()' failed.    # NOQA
@@ -199,6 +200,7 @@ def test_attn(
         kernel=real_kernel,
         enable_tma=enable_tma,
         scaling_seqlen=scaling_seqlen,
+        fp8_quant_mode=fp8_quant_mode,
     )
 
     torch.testing.assert_close(
@@ -584,6 +586,53 @@ def test_attn_cutlass(self, *args, **kwargs) -> None:
                 real_kernel=Kernel.CUTLASS,
             )
 
+    # Samples every FP8 quant_mode; only SM90 (Hopper) supports all of them.
+    # SM120 (Blackwell RTX) is mode=2 fwd-only; SM80/SM100 have no FP8.
+    @mark_ci_scope("h20")
+    @unittest.skipIf(*gpu_unavailable)
+    @unittest.skipIf(
+        not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 9,
+        "all-FP8-mode test requires SM90 (Hopper)",
+    )
+    # pyre-ignore
+    @given(
+        batch_size=st.integers(4, 8),
+        heads=st.integers(1, 4),
+        max_uih_len=st.sampled_from([20, 100, 128]),
+        max_targets=st.sampled_from([20, 512]),
+        attn_dim=st.sampled_from([64, 128]),
+        causal=st.sampled_from([True]),
+        has_multiple_targets=st.sampled_from([True, False]),
+        dtype=st.sampled_from(get_test_dtypes([torch.bfloat16])),
+        has_max_attn_len=st.sampled_from([False]),
+        contextual_seq_len=st.sampled_from([0, 10]),
+        scaling_seqlen=st.sampled_from([-1, 2048]),
+        # mode=1 (two-direction) excluded: the wheel's vt TMA descriptor
+        # init fails for some shapes (e.g. batch=6, heads=4, max_uih=100,
+        # attn_dim=128). Revisit once the wheel is fixed upstream.
+        fp8_quant_mode=st.sampled_from([0, 2, 3, 4, 5]),
+    )
+    @settings(
+        verbosity=Verbosity.verbose,
+        max_examples=20,
+        deadline=None,
+    )
+    # pyre-ignore[2]
+    def test_attn_fp8_cutlass(self, *args, **kwargs) -> None:
+        # FP8 e4m3 vs bf16 ref; tolerances relaxed for quantization loss.
+        hidden_dim = kwargs.pop("attn_dim")
+        test_attn(
+            *args,
+            **kwargs,
+            attn_dim=hidden_dim,
+            hidden_dim=hidden_dim,
+            test_backward=True,
+            ref_kernel=Kernel.PYTORCH,
+            real_kernel=Kernel.CUTLASS,
+            atol=1e-1,
+            rtol=1e-1,
+        )
+
     # NOTE: no ``test_delta_attn_cutlass`` — ``delta_hstu_mha`` has no
     # CUTLASS implementation and falls back to Triton internally. The
     # delta/cached path is already covered by ``test_delta_attn_triton``.
@@ -603,7 +652,7 @@ def test_attn_cutlass(self, *args, **kwargs) -> None:
         sla_k2=st.sampled_from([0, 4, 8]),
         has_multiple_targets=st.sampled_from([True, False]),
         contextual_seq_len=st.sampled_from([0, 4]),
-        # Sample both bf16 and fp16 -- fp8 is deferred to ultra-hstu-fp8.
+        # Sample both bf16 and fp16; FP8 is covered by test_attn_fp8_cutlass.
         dtype=st.sampled_from(get_test_dtypes([torch.bfloat16, torch.float16])),
     )
     @settings(