diff --git a/docs/source/models/dlrm_hstu.md b/docs/source/models/dlrm_hstu.md index 27a01c2b2..558bf5e1b 100644 --- a/docs/source/models/dlrm_hstu.md +++ b/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 配置方式 diff --git a/requirements/extra.txt b/requirements/extra.txt index aa46b5fed..bfcfc6a2d 100644 --- a/requirements/extra.txt +++ b/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 diff --git a/tzrec/modules/gr/stu.py b/tzrec/modules/gr/stu.py index 5bee0b002..6965c7fcb 100644 --- a/tzrec/modules/gr/stu.py +++ b/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( diff --git a/tzrec/ops/_cuda/cutlass_hstu_attention.py b/tzrec/ops/_cuda/cutlass_hstu_attention.py index 37f3f22d0..5d760b9ef 100644 --- a/tzrec/ops/_cuda/cutlass_hstu_attention.py +++ b/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, ) diff --git a/tzrec/ops/hstu_attention.py b/tzrec/ops/hstu_attention.py index 4e0bb2786..6287eadeb 100644 --- a/tzrec/ops/hstu_attention.py +++ b/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: diff --git a/tzrec/ops/hstu_attention_test.py b/tzrec/ops/hstu_attention_test.py index 6a653a464..8a484abb9 100644 --- a/tzrec/ops/hstu_attention_test.py +++ b/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::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( diff --git a/tzrec/ops/hstu_compute.py b/tzrec/ops/hstu_compute.py index b53b48ca3..4f01aeee8 100644 --- a/tzrec/ops/hstu_compute.py +++ b/tzrec/ops/hstu_compute.py @@ -323,6 +323,7 @@ def hstu_preprocess_and_attention( enable_tma: bool = False, attn_func: Optional[torch.Tensor] = None, scaling_seqlen: int = -1, + fp8_quant_mode: int = -1, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]: if not is_fx_tracing(): torch._assert(max_seq_len > 0, "max_seq_len must be larger than 0") @@ -350,6 +351,12 @@ def hstu_preprocess_and_attention( "split the call into separate preprocess and hstu_mha " "(prefill=True)." ) + if fp8_quant_mode >= 0: + raise ValueError( + f"fp8_quant_mode={fp8_quant_mode} (FP8) is not supported on " + "the fused Triton preprocess+attention path; FP8 attention " + "requires kernel=Kernel.CUTLASS." + ) from tzrec.ops._triton.triton_hstu_preprocess_and_attention import ( triton_hstu_preprocess_and_attention, ) @@ -417,5 +424,6 @@ def hstu_preprocess_and_attention( kernel=kernel, attn_func=attn_func, scaling_seqlen=scaling_seqlen, + fp8_quant_mode=fp8_quant_mode, ).view(-1, hidden_dim * num_heads) return u, attn_output, k, v diff --git a/tzrec/protos/module.proto b/tzrec/protos/module.proto index 2b5b5ec2d..5492d1974 100644 --- a/tzrec/protos/module.proto +++ b/tzrec/protos/module.proto @@ -272,6 +272,12 @@ message STU { // attention output scaling divisor (denominator of the SiLU(QK)/N term). // Sentinel: < 0 (default) = use runtime max_seq_len. optional int32 scaling_seqlen = 16 [default = -1]; + // FP8 attention quant mode for the CUTLASS kernel. + // -1 (default) = off (bf16/fp16); 0..5 select FP8 modes (0 per-tensor, + // 1 two-direction, 2 per-block, 3 per-head, 4 per-batch, 5 global). + // Requires model.kernel = CUTLASS on SM90 (Hopper, all modes, fwd+bwd) + // or SM120 (Blackwell RTX, quant_mode=2 only, forward-only). + optional int32 fp8_quant_mode = 17 [default = -1]; } message GRPositionalEncoder { diff --git a/tzrec/tests/rank_integration_test.py b/tzrec/tests/rank_integration_test.py index a163d00fb..66dc8ed49 100644 --- a/tzrec/tests/rank_integration_test.py +++ b/tzrec/tests/rank_integration_test.py @@ -1098,10 +1098,19 @@ def test_rank_dlrm_hstu_cutlass_train_eval_export(self): @unittest.skipIf(*cutlass_hstu_unavailable) @unittest.skipIf(*gpu_unavailable) def test_rank_ultra_hstu_cutlass_train_eval_export(self): - self.success = utils.test_train_eval( - "tzrec/tests/configs/ultra_hstu_cutlass_kuairand_1k.config", - self.test_dir, - ) + # Enable per-block FP8 attention (quant_mode=2) for the H20 run -- + # exercises train + eval + AOTI export + predict end-to-end on the + # CUTLASS FP8 path. Export routes through the opaque + # tzrec::cutlass_hstu_fp8_fwd custom op (see cutlass_hstu_mha) so the + # wheel's data-dependent Python quantizer never enters the graph. + pc = config_util.load_pipeline_config( + "tzrec/tests/configs/ultra_hstu_cutlass_kuairand_1k.config" + ) + for hstu in pc.model_config.ultra_hstu.hstu: + hstu.stu.fp8_quant_mode = 2 + fp8_config = os.path.join(self.test_dir, "ultra_hstu_cutlass_fp8.config") + config_util.save_message(pc, fp8_config) + self.success = utils.test_train_eval(fp8_config, self.test_dir) if self.success: self.success = utils.test_eval( os.path.join(self.test_dir, "pipeline.config"), self.test_dir diff --git a/tzrec/version.py b/tzrec/version.py index 4ac3e6f3c..8bb46cf62 100644 --- a/tzrec/version.py +++ b/tzrec/version.py @@ -9,4 +9,4 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "1.2.15" +__version__ = "1.2.16"