diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index 1d45d1919af..42f65d61eab 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -945,3 +945,57 @@ def test_fp8_current_scaling_with_layernorm_linear_module( recipe1_golden_tensors=None, recipe2_golden_tensors=fp8_zero_tolerance_tensor_dumps_recipe2, ) + + +@pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) +class TestFP8CurrentScalingTritonQuantize: + """The two-phase Triton quantization must be bitwise equal to the C++ kernels.""" + + @staticmethod + def _quantize(x, fp8_dtype, pow2, eps, use_triton): + import transformer_engine.pytorch.tensor.float8_tensor as f8 + + quantizer = te.Float8CurrentScalingQuantizer( + fp8_dtype=fp8_dtype, + device=torch.device("cuda"), + rowwise=True, + columnwise=False, + force_pow_2_scales=pow2, + amax_epsilon=eps, + ) + old = f8._TRITON_CURRENT_SCALING, f8._TRITON_CURRENT_SCALING_MIN_NUMEL + f8._TRITON_CURRENT_SCALING, f8._TRITON_CURRENT_SCALING_MIN_NUMEL = use_triton, 1 + try: + return quantizer(x) + finally: + f8._TRITON_CURRENT_SCALING, f8._TRITON_CURRENT_SCALING_MIN_NUMEL = old + + @pytest.mark.parametrize("fp8_dtype", [tex.DType.kFloat8E4M3, tex.DType.kFloat8E5M2]) + @pytest.mark.parametrize("pow2", [False, True]) + @pytest.mark.parametrize("eps", [0.0, 1e-2]) + @pytest.mark.parametrize("numel", [1 << 20, (1 << 20) + 37, 3 * 8192 * 5 + 1]) + @pytest.mark.parametrize("special", [None, "zeros", "inf", "nan"]) + def test_triton_matches_cpp(self, fp8_dtype, pow2, eps, numel, special): + torch.manual_seed(0) + x = torch.randn(numel, dtype=torch.bfloat16, device="cuda") * 3 + if special == "zeros": + x.zero_() + elif special == "inf": + x[numel // 2] = float("inf") + elif special == "nan": + x[numel // 3] = float("nan") + ref = self._quantize(x, fp8_dtype, pow2, eps, use_triton=False) + out = self._quantize(x, fp8_dtype, pow2, eps, use_triton=True) + torch.testing.assert_close(out._scale_inv, ref._scale_inv, rtol=0, atol=0) + assert torch.equal(out._data.view(torch.uint8), ref._data.view(torch.uint8)) + + def test_repeated_calls_reuse_workspace(self): + x = torch.randn(1 << 22, dtype=torch.bfloat16, device="cuda") + first = self._quantize(x, tex.DType.kFloat8E5M2, False, 0.0, True) + for scale in (1.0, 100.0, 0.01): + y = x * scale + ref = self._quantize(y, tex.DType.kFloat8E5M2, False, 0.0, False) + out = self._quantize(y, tex.DType.kFloat8E5M2, False, 0.0, True) + assert torch.equal(out._scale_inv, ref._scale_inv) + assert torch.equal(out._data.view(torch.uint8), ref._data.view(torch.uint8)) + assert first._data is not None diff --git a/tests/pytorch/test_fused_rope.py b/tests/pytorch/test_fused_rope.py index 5cc4fc0f8a0..456407dfce3 100644 --- a/tests/pytorch/test_fused_rope.py +++ b/tests/pytorch/test_fused_rope.py @@ -497,3 +497,92 @@ def test_rotary_position_embedding_forward_with_autocast_gives_same_result_as_wi atol=1e-8, rtol=1e-8, ) + + +def _fp8_current_scaling_available() -> bool: + import transformer_engine.pytorch as te + + return te.is_fp8_available() + + +@pytest.mark.skipif(not _fp8_current_scaling_available(), reason="FP8 is not available") +@pytest.mark.parametrize("seq_length", [128, 1024, 2048 + 8]) +@pytest.mark.parametrize("batch_size", [1, 2]) +@pytest.mark.parametrize("groups, q_per_group", [(8, 4), (4, 1), (2, 8)]) +@pytest.mark.parametrize("pow2", [False, True]) +def test_fused_rope_fp8_quantize_matches_unfused( + seq_length: int, batch_size: int, groups: int, q_per_group: int, pow2: bool +) -> None: + """Deferred RoPE + FP8 quantize is bitwise equal to TE's RoPE followed by quantization.""" + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from transformer_engine.pytorch.attention import deferred_qkv_rope + + head_dim = 128 + split = [q_per_group * head_dim, head_dim, head_dim] + device = torch.device("cuda:0") + torch.manual_seed(0) + qkv = ( + torch.randn(seq_length, batch_size, groups, sum(split), dtype=torch.bfloat16, device=device) + * 4 + ) + emb = RotaryPositionEmbedding(head_dim, 1.0)(seq_length + 3) + + q_ref, k_ref, v_ref = apply_fused_qkv_rotary_pos_emb(qkv, emb, emb, split) + quantizer = te.Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device=device, + rowwise=True, + columnwise=False, + force_pow_2_scales=pow2, + ) + ref_data = torch.cat([t.reshape(-1) for t in (q_ref, k_ref, v_ref)]) + ref = quantizer(ref_data) + + q, k, v = apply_fused_qkv_rotary_pos_emb(qkv, emb, emb, split, defer_for_fp8_attention=True) + assert deferred_qkv_rope.get_state(q) is not None + q_fp8, k_fp8, v_fp8 = deferred_qkv_rope.quantize_deferred(q, k, v, quantizer) + + torch.testing.assert_close(q_fp8._scale_inv, ref._scale_inv, rtol=0, atol=0) + out = torch.cat([t._data.reshape(-1) for t in (q_fp8, k_fp8, v_fp8)]) + assert torch.equal(out.view(torch.uint8), ref._data.view(torch.uint8)) + assert q_fp8.shape == q_ref.shape and k_fp8.shape == k_ref.shape + + +@pytest.mark.parametrize("seq_length", [128, 1024]) +def test_deferred_rope_materialize_and_backward(seq_length: int) -> None: + """Placeholders filled by materialize() and their gradient match the regular fused RoPE.""" + from transformer_engine.pytorch.attention import deferred_qkv_rope + + head_dim, groups, q_per_group, batch_size = 128, 8, 4, 2 + split = [q_per_group * head_dim, head_dim, head_dim] + device = torch.device("cuda:0") + torch.manual_seed(0) + qkv = torch.randn( + seq_length, batch_size, groups, sum(split), dtype=torch.bfloat16, device=device + ).requires_grad_() + emb = RotaryPositionEmbedding(head_dim, 1.0)(seq_length) + + ref = apply_fused_qkv_rotary_pos_emb(qkv, emb, emb, split) + loss_ref = sum(t.float().sum() * 2 for t in ref) + (grad_ref,) = torch.autograd.grad(loss_ref, qkv) + + out = apply_fused_qkv_rotary_pos_emb(qkv, emb, emb, split, defer_for_fp8_attention=True) + assert deferred_qkv_rope.get_state(out[0]) is not None + deferred_qkv_rope.materialize(*out) + for a, b in zip(out, ref): + torch.testing.assert_close(a, b, rtol=0, atol=0) + loss = sum(t.float().sum() * 2 for t in out) + (grad,) = torch.autograd.grad(loss, qkv) + torch.testing.assert_close(grad, grad_ref, rtol=0, atol=0) + + +def test_deferred_rope_unsupported_falls_back() -> None: + """Unsupported inputs return ordinary tensors instead of placeholders.""" + from transformer_engine.pytorch.attention import deferred_qkv_rope + + split = [4 * 64, 64, 64] # head dim != 128 + qkv = torch.randn(16, 2, 8, sum(split), dtype=torch.bfloat16, device="cuda:0") + emb = RotaryPositionEmbedding(64, 1.0)(16) + q, _, _ = apply_fused_qkv_rotary_pos_emb(qkv, emb, emb, split, defer_for_fp8_attention=True) + assert deferred_qkv_rope.get_state(q) is None diff --git a/transformer_engine/common/triton/current_scaling.py b/transformer_engine/common/triton/current_scaling.py new file mode 100644 index 00000000000..7b732b924c8 --- /dev/null +++ b/transformer_engine/common/triton/current_scaling.py @@ -0,0 +1,69 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Two-phase FP8 current-scaling quantization kernels written with OpenAI Triton.""" + +import triton +import triton.language as tl + +try: + from triton.language.extra import libdevice +except ImportError: # older Triton + from triton.language.extra.cuda import libdevice + + +@triton.jit +def _finalize_scale(ws_ptr, scale_inv_ptr, n_prog, eps, FP8MAX: tl.constexpr, POW2: tl.constexpr): + """Executed by every CTA after its amax atomic; the last one turns amax into the scale. + + Mirrors TE's compute_scale_from_amax. The workspace ``ws = [amax, counter, scale]`` is + handed back zeroed so no memset launch is needed (CUDA-graph safe). + """ + cnt = tl.atomic_add(ws_ptr + 1, 1.0) + if cnt == n_prog - 1: + amax = tl.atomic_xchg(ws_ptr, 0.0) + tl.atomic_xchg(ws_ptr + 1, 0.0) + bad = (amax == float("inf")) | (amax != amax) + amax = tl.where(amax < eps, eps, amax) + bad = bad | (amax == 0.0) + scale = libdevice.div_rn(FP8MAX, amax) + scale = tl.where(scale == float("inf"), 3.4028234663852886e38, scale) + if POW2: + scale = (scale.to(tl.uint32, bitcast=True) & 0xFF800000).to(tl.float32, bitcast=True) + scale = tl.where(bad, 1.0, scale) + tl.store(ws_ptr + 2, scale) + tl.store(scale_inv_ptr, libdevice.rcp_rn(scale)) + + +@triton.jit +def current_scaling_quantize_kernel( + x_ptr, + out_ptr, + ws_ptr, + scale_inv_ptr, + n, + eps, + BLOCK: tl.constexpr, + CAST: tl.constexpr, + FP8MAX: tl.constexpr, + POW2: tl.constexpr, +): + """CAST=False: persistent amax pass, one atomic per CTA, last CTA computes the scale. + CAST=True: one block per CTA, in reverse order so the amax pass's tail is still in L2.""" + n_prog = tl.num_programs(0) + pid = tl.program_id(0) + if CAST: + offs = (n_prog - 1 - pid).to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + m = offs < n + x = tl.load(x_ptr + offs, mask=m, other=0.0).to(tl.float32) + scale = tl.load(ws_ptr + 2) + tl.store(out_ptr + offs, (x * scale).to(out_ptr.dtype.element_ty), mask=m) + else: + amax = tl.zeros([BLOCK], tl.float32) + for blk in tl.range(pid, tl.cdiv(n, BLOCK), n_prog): + offs = blk.to(tl.int64) * BLOCK + tl.arange(0, BLOCK) + x = tl.load(x_ptr + offs, mask=offs < n, other=0.0).to(tl.float32) + amax = tl.maximum(amax, tl.abs(x)) + tl.atomic_max(ws_ptr, tl.max(amax, axis=0)) + _finalize_scale(ws_ptr, scale_inv_ptr, n_prog, eps, FP8MAX, POW2) diff --git a/transformer_engine/common/triton/rope_fp8.py b/transformer_engine/common/triton/rope_fp8.py new file mode 100644 index 00000000000..54ed3a16465 --- /dev/null +++ b/transformer_engine/common/triton/rope_fp8.py @@ -0,0 +1,169 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Fused RoPE + FP8 current-scaling quantization of Q/K/V, written with OpenAI Triton. + +Input is the QKV GEMM output ``[s, b, g, (hq + 2) * d]``. Two passes are needed because current +scaling requires the global amax before the first cast: + +* amax pass: persistent; rotates Q/K (rounded to bf16 like TE's RoPE output) and reduces + amax(|Q|, |K|, |V|) with one atomic per CTA; the last CTA turns the amax into the scale. +* cast pass: rotates again and writes e4m3/e5m2 into the packed ``[Q | K | V]`` payload, visiting + tiles in reverse order so the tail of the amax pass is still in L2. + +Results are bitwise identical to ``tex.fused_qkv_rope_forward`` followed by current-scaling +quantization. +""" + +import triton +import triton.language as tl + +from .current_scaling import _finalize_scale + +try: + from triton.language.extra import libdevice +except ImportError: # older Triton + from triton.language.extra.cuda import libdevice + + +@triton.jit +def rope_trig_table_kernel(f_ptr, cos_ptr, sin_ptr, n, BLOCK: tl.constexpr): + """cos/sin of the RoPE angles with the libdevice functions TE's sincosf resolves to, + evaluated once per angle instead of once per element of every head.""" + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = offs < n + f = tl.load(f_ptr + offs, mask=m, other=0.0) + tl.store(cos_ptr + offs, libdevice.cos(f), mask=m) + tl.store(sin_ptr + offs, libdevice.sin(f), mask=m) + + +@triton.jit +def _rot(x, c, xr, s): + # fma(x, c, xr * s): matches nvcc's contraction in TE's RoPE kernels + return tl.fma(x, c, xr * s) + + +@triton.jit +def _rot_fwd(lo, hi, c_lo, s_lo, c_hi, s_hi): + # TE fused_qkv_rope_block_forward: out[d=h] = x[d]cos[d] + x[d-h]sin[d]; bf16 output + y_lo = _rot(lo, c_lo, -hi, s_lo).to(tl.bfloat16).to(tl.float32) + y_hi = _rot(hi, c_hi, lo, s_hi).to(tl.bfloat16).to(tl.float32) + return y_lo, y_hi + + +@triton.jit +def _rope_qkv_token( + x_ptr, + cos_ptr, + sin_ptr, + out_ptr, + sc, + tok, + n_tok, + B: tl.constexpr, + G: tl.constexpr, + HQ: tl.constexpr, + D: tl.constexpr, + CAST: tl.constexpr, +): + """Rotate Q/K of one (s, b) token for all heads. CAST=False returns the token's amax; + CAST=True writes the scaled FP8 values into the packed [Q | K | V] payload.""" + W: tl.constexpr = (HQ + 2) * D + HD: tl.constexpr = D // 2 + pos = tok // B + g = tl.arange(0, G) + jq = tl.arange(0, HQ) + dl = tl.arange(0, HD) + d = tl.arange(0, D) + tb = pos * D + dl + c_lo = tl.load(cos_ptr + tb) + s_lo = tl.load(sin_ptr + tb) + c_hi = tl.load(cos_ptr + tb + HD) + s_hi = tl.load(sin_ptr + tb + HD) + # Q: [g, jq, dl]; K: [g, dl]; V: [g, d] + row = tok * (G * W) + g * W + q_src = x_ptr + row[:, None, None] + jq[None, :, None] * D + dl[None, None, :] + k_src = x_ptr + row[:, None] + HQ * D + dl[None, :] + v_src = x_ptr + row[:, None] + (HQ + 1) * D + d[None, :] + q_lo, q_hi = _rot_fwd( + tl.load(q_src).to(tl.float32), + tl.load(q_src + HD).to(tl.float32), + c_lo[None, None, :], + s_lo[None, None, :], + c_hi[None, None, :], + s_hi[None, None, :], + ) + k_lo, k_hi = _rot_fwd( + tl.load(k_src).to(tl.float32), + tl.load(k_src + HD).to(tl.float32), + c_lo[None, :], + s_lo[None, :], + c_hi[None, :], + s_hi[None, :], + ) + v = tl.load(v_src).to(tl.float32) + if CAST: + nq = n_tok * (G * HQ * D) + nk = n_tok * (G * D) + q_dst = ( + out_ptr + + tok * (G * HQ * D) + + g[:, None, None] * (HQ * D) + + jq[None, :, None] * D + + dl[None, None, :] + ) + k_dst = out_ptr + nq + tok * (G * D) + g[:, None] * D + dl[None, :] + v_dst = out_ptr + nq + nk + tok * (G * D) + g[:, None] * D + d[None, :] + ety = out_ptr.dtype.element_ty + tl.store(q_dst, (q_lo * sc).to(ety)) + tl.store(q_dst + HD, (q_hi * sc).to(ety)) + tl.store(k_dst, (k_lo * sc).to(ety)) + tl.store(k_dst + HD, (k_hi * sc).to(ety)) + tl.store(v_dst, (v * sc).to(ety)) + return 0.0 + else: + aq = tl.max(tl.max(tl.max(tl.maximum(tl.abs(q_lo), tl.abs(q_hi)), 2), 1), 0) + ak = tl.max(tl.max(tl.maximum(tl.abs(k_lo), tl.abs(k_hi)), 1), 0) + av = tl.max(tl.max(tl.abs(v), 1), 0) + return tl.maximum(tl.maximum(aq, ak), av) + + +@triton.jit +def rope_qkv_quantize_kernel( + x_ptr, + cos_ptr, + sin_ptr, + out_ptr, + ws_ptr, + scale_inv_ptr, + n_tok, + eps, + B: tl.constexpr, + G: tl.constexpr, + HQ: tl.constexpr, + D: tl.constexpr, + CAST: tl.constexpr, + FP8MAX: tl.constexpr, + POW2: tl.constexpr, + NUM_STAGES: tl.constexpr = 1, +): + """CAST=False: persistent amax pass (one atomic per CTA, last CTA computes the scale). + CAST=True: one token per CTA, in the reverse order of the amax pass.""" + n_prog = tl.num_programs(0) + pid = tl.program_id(0) + if CAST: + sc = tl.load(ws_ptr + 2) + _rope_qkv_token( + x_ptr, cos_ptr, sin_ptr, out_ptr, sc, n_prog - 1 - pid, n_tok, B, G, HQ, D, True + ) + else: + amax = 0.0 + for tok in tl.range(pid, n_tok, n_prog, num_stages=NUM_STAGES): + a = _rope_qkv_token( + x_ptr, cos_ptr, sin_ptr, out_ptr, 1.0, tok, n_tok, B, G, HQ, D, False + ) + amax = tl.maximum(amax, a) + tl.atomic_max(ws_ptr, amax) + _finalize_scale(ws_ptr, scale_inv_ptr, n_prog, eps, FP8MAX, POW2) diff --git a/transformer_engine/pytorch/attention/deferred_qkv_rope.py b/transformer_engine/pytorch/attention/deferred_qkv_rope.py new file mode 100644 index 00000000000..057f8221f02 --- /dev/null +++ b/transformer_engine/pytorch/attention/deferred_qkv_rope.py @@ -0,0 +1,228 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Deferred fused-QKV RoPE for FP8 current-scaling attention. + +``apply_fused_qkv_rotary_pos_emb(..., defer_for_fp8_attention=True)`` returns *placeholder* Q/K/V +tensors tagged with the QKV GEMM output and the RoPE angles instead of running RoPE. If the +DotProductAttention that receives them runs FP8 current-scaling fused attention, RoPE and the +FP8 quantization of Q/K/V happen in one fused kernel straight from the GEMM output and the +placeholders are never read. Otherwise DotProductAttention fills them with TE's regular RoPE +before doing anything else, so results do not depend on which path is taken. +""" + +import os +from typing import List, Optional, Tuple + +import torch + +import transformer_engine_torch as tex + +_ENABLED = os.environ.get("NVTE_FUSED_QKV_ROPE_FP8", "1") == "1" +_HEAD_DIM = 128 + + +class _DeferredRoPEState: + """What the deferred RoPE needs, carried from the RoPE call to the quantize call.""" + + __slots__ = ("mixed", "freqs", "split", "cos_sin", "consumed", "materialized") + + def __init__(self, mixed: torch.Tensor, freqs: torch.Tensor, split: List[int]): + self.mixed = mixed # released once quantized or materialized + self.freqs = freqs + self.split = list(split) + self.cos_sin = None + self.consumed = False + self.materialized = False + + +class _DeferredQKVRoPEFunc(torch.autograd.Function): + """Placeholders whose backward is TE's fused QKV RoPE backward.""" + + @staticmethod + def forward(ctx, mixed, state): + s, b, g, _ = mixed.shape + d = state.split[1] + hq = state.split[0] // d + kwargs = dict(dtype=mixed.dtype, device=mixed.device) + q = torch.empty(s, b, g * hq, d, **kwargs) + k = torch.empty(s, b, g, d, **kwargs) + v = torch.empty(s, b, g, d, **kwargs) + ctx.freqs = state.freqs + ctx.split = state.split + return q, k, v + + @staticmethod + def backward(ctx, grad_q, grad_k, grad_v): + # pylint: disable=import-outside-toplevel + from .rope import QKVFormat + + grad_input = tex.fused_qkv_rope_backward( + grad_q.contiguous(), + grad_k.contiguous(), + grad_v.contiguous(), + ctx.freqs, + ctx.freqs, + ctx.split, + QKVFormat["sbhd"], + False, + 1, + 0, + ) + return grad_input, None + + +def get_state(q: torch.Tensor) -> Optional[_DeferredRoPEState]: + """Deferred-RoPE state of a placeholder, or None for an ordinary tensor.""" + return getattr(q, "_deferred_qkv_rope", None) + + +def defer_fused_qkv_rope( + qkv: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + qkv_split_arg_list: List[int], + start_positions, + tensor_format: str, + interleaved: bool, + cp_size: int, +) -> Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """Return tagged placeholders for the deferred RoPE, or None if the inputs are unsupported.""" + split = list(qkv_split_arg_list) + if not ( + _ENABLED + and start_positions is None + and tensor_format == "sbhd" + and not interleaved + and cp_size == 1 + and qkv.dim() == 4 + and qkv.is_cuda + and qkv.dtype == torch.bfloat16 + and qkv.is_contiguous() + and len(split) == 3 + and split[1] == split[2] == _HEAD_DIM + and split[0] % split[1] == 0 + and qkv.shape[3] == sum(split) + and q_freqs.data_ptr() == k_freqs.data_ptr() + and q_freqs.shape == k_freqs.shape + and q_freqs.stride() == k_freqs.stride() + and q_freqs.dim() == 4 + and q_freqs.shape[1] == 1 + and q_freqs.shape[2] == 1 + and q_freqs.shape[3] == split[1] + and q_freqs.shape[0] >= qkv.shape[0] + and q_freqs.is_cuda + ): + return None + freqs = q_freqs if q_freqs.dtype == torch.float32 else q_freqs.float() + if not freqs.is_contiguous(): + freqs = freqs.contiguous() + state = _DeferredRoPEState(qkv, freqs, split) + q, k, v = _DeferredQKVRoPEFunc.apply(qkv, state) + for t in (q, k, v): + t._deferred_qkv_rope = state + return q, k, v + + +def materialize(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> None: + """Run TE's regular RoPE into the placeholders (no-op for ordinary tensors).""" + state = get_state(q) + if state is None or state.materialized: + return + if state.mixed is None: + raise RuntimeError( + "Deferred fused-QKV RoPE placeholders were already consumed by the fused FP8 " + "quantization and cannot be reused (e.g. attention recompute). Disable " + "defer_for_fp8_attention / NVTE_FUSED_QKV_ROPE_FP8 for this configuration." + ) + with torch.no_grad(): + q_rot, k_rot, v_rot = tex.fused_qkv_rope_forward( + state.mixed, + state.freqs, + state.freqs, + None, + state.split, + _sbhd(), + False, + 1, + 0, + ) + q.copy_(q_rot) + k.copy_(k_rot) + v.copy_(v_rot) + state.consumed = state.materialized = True + state.mixed = None + + +def _sbhd(): + # pylint: disable=import-outside-toplevel + from .rope import QKVFormat + + return QKVFormat["sbhd"] + + +def can_fuse_with_fp8_attention( + q: torch.Tensor, qkv_quantizer, qkv_layout: str, is_bwd_fp8: bool +) -> bool: + """Whether FP8 fused attention with this QKV quantizer can consume the deferred RoPE.""" + # pylint: disable=import-outside-toplevel + from ..tensor.float8_tensor import Float8CurrentScalingQuantizer + from ..constants import DType + + state = get_state(q) + return ( + state is not None + and not state.consumed + and is_bwd_fp8 + and type(qkv_quantizer) is Float8CurrentScalingQuantizer + and not qkv_quantizer.with_amax_reduction + and qkv_quantizer.dtype == DType.kFloat8E4M3 + and qkv_layout == "sbhd_sbhd_sbhd" + ) + + +def quantize_deferred(q, k, v, qkv_quantizer): + """Fused RoPE + FP8 quantization of deferred placeholders. + + Returns ``(q_fp8, k_fp8, v_fp8)`` Float8Tensor views over one packed payload, or None if q is + not a placeholder (or the quantizer's output layout is unsupported, in which case the + placeholders are materialized with TE's RoPE and the regular quantization applies). + """ + # pylint: disable=import-outside-toplevel + from ..tensor.float8_tensor import Float8Tensor + from ..triton.rope_fp8 import rope_trig_tables, rope_qkv_quantize + + state = get_state(q) + if state is None: + return None + if state.consumed: + raise RuntimeError("Deferred fused-QKV RoPE placeholders were already consumed.") + d = state.split[1] + hq = state.split[0] // d + s, b, g, _ = state.mixed.shape + nq, nk = s * b * g * hq * d, s * b * g * d + qkv_fp8 = qkv_quantizer.make_empty((nq + 2 * nk,), dtype=q.dtype, device=q.device) + if qkv_fp8._data is None or qkv_fp8._transpose is not None or qkv_fp8._scale_inv.numel() != 1: + materialize(q, k, v) + return None + data = qkv_fp8._data + cos, sin = rope_trig_tables(state.freqs) + rope_qkv_quantize( + state.mixed, + cos, + sin, + data.view(torch.float8_e4m3fn), + qkv_fp8._scale_inv, + hq, + eps=float(qkv_quantizer.amax_epsilon), + force_pow_2_scales=bool(qkv_quantizer.force_pow_2_scales), + ) + views = ( + data[:nq].view(q.shape), + data[nq : nq + nk].view(k.shape), + data[nq + nk :].view(v.shape), + ) + state.consumed = True + state.mixed = None + return tuple(Float8Tensor.make_like(qkv_fp8, data=x, dtype=q.dtype) for x in views) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index c3503d0f1bf..36d4a4abb7f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -85,6 +85,7 @@ # Import attention utils import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils +from transformer_engine.pytorch.attention import deferred_qkv_rope from transformer_engine.pytorch.attention.dot_product_attention.utils import ( FlashAttentionUtils as fa_utils, combine_and_quantize, @@ -1793,16 +1794,22 @@ def _fused_attn_forward_impl( if fp8_recipe.mxfp8(): qkv_scale_inv_format = "bhsd" # Same as what combine_and_quantize would give else: - q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( - qkv_layout, - q, - k, - v, - QKV_quantizer, - used_in_backward=is_training, - combined_qkv=args.packed_qkv, - combined_kv=args.packed_kv, - ) + # RoPE placeholders from apply_fused_qkv_rotary_pos_emb(defer_for_fp8_attention=True) + # are rotated and quantized together; None for ordinary tensors + deferred_qkv = deferred_qkv_rope.quantize_deferred(q, k, v, QKV_quantizer) + if deferred_qkv is not None: + q_fp8, k_fp8, v_fp8 = deferred_qkv + else: + q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( + qkv_layout, + q, + k, + v, + QKV_quantizer, + used_in_backward=is_training, + combined_qkv=args.packed_qkv, + combined_kv=args.packed_kv, + ) # print quantizers print_quantizers( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 9a3c8a44dbf..6e1b66172d1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -57,6 +57,7 @@ from transformer_engine.pytorch.utils import SplitAlongDim from transformer_engine.pytorch.graph import is_graph_capturing from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.attention import deferred_qkv_rope # Import attention utils import transformer_engine.pytorch.attention.dot_product_attention.utils as dpa_utils @@ -2330,6 +2331,7 @@ def forward( qkv_interleave_dim, inference_params, ) + deferred_rope_state = deferred_qkv_rope.get_state(query_layer) with self.prepare_forward_ctx( query_layer, @@ -2921,6 +2923,36 @@ def forward( " disabling all backends." ) + # Q/K/V returned by apply_fused_qkv_rotary_pos_emb(defer_for_fp8_attention=True) hold no + # data yet: either the fused FP8 attention path quantizes them from the QKV GEMM output, + # or TE's regular RoPE is applied into them here. + if deferred_rope_state is not None: + if deferred_qkv_rope.get_state(query_layer) is not deferred_rope_state: + raise RuntimeError( + "Q/K/V from apply_fused_qkv_rotary_pos_emb(defer_for_fp8_attention=True) " + "were replaced before attention; pass them to DotProductAttention as is." + ) + fp8_dpa = self.fp8 and self.fp8_meta["recipe"].fp8_dpa + can_fuse = ( + fp8_dpa + and use_fused_attention + and not ( + checkpoint_core_attention + or context_parallel + or bf16_backward + or score_mod is not None + or inference_params is not None + ) + and deferred_qkv_rope.can_fuse_with_fp8_attention( + query_layer, + dpa_utils.get_attention_quantizers(fp8_dpa, self.quantizers)[0], + qkv_layout, + bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))), + ) + ) + if not can_fuse: + deferred_qkv_rope.materialize(query_layer, key_layer, value_layer) + # run attention softmax_offset = ( self.softmax_offset.reshape(1, -1, 1, 1).to(torch.float32) diff --git a/transformer_engine/pytorch/attention/rope.py b/transformer_engine/pytorch/attention/rope.py index 83369547f1f..4c2f46ddf53 100644 --- a/transformer_engine/pytorch/attention/rope.py +++ b/transformer_engine/pytorch/attention/rope.py @@ -474,6 +474,7 @@ def apply_fused_qkv_rotary_pos_emb( cu_seqlens: Union[torch.Tensor, None] = None, # pylint: disable=unused-argument cp_size: int = 1, cp_rank: int = 0, + defer_for_fp8_attention: bool = False, ) -> torch.Tensor: """ Apply rotary positional embedding tensor to the input qkv tensor. @@ -520,6 +521,14 @@ def apply_fused_qkv_rotary_pos_emb( Context parallel world size. cp_rank : int, default = 0. Context parallel rank. + defer_for_fp8_attention : bool, default = False + Return placeholder Q, K, V that must be passed straight to a + :class:`DotProductAttention`. With FP8 current-scaling fused attention, RoPE and the + FP8 quantization of Q, K, V then run as one fused kernel from ``qkv``; in every other + case DotProductAttention applies the regular fused RoPE to them first. Falls back to + the regular RoPE for unsupported inputs (bf16 sbhd, head dim 128, ``cp_size == 1``, + no ``start_positions``, non-interleaved). The placeholders cannot be reused after + attention consumed them (e.g. attention recompute). """ # `start_positions` is only supported for `cp_size=1` and inference. @@ -529,6 +538,23 @@ def apply_fused_qkv_rotary_pos_emb( assert tensor_format != "thd", "'thd' tensor_format not supported currently." + if defer_for_fp8_attention: + # pylint: disable=import-outside-toplevel + from .deferred_qkv_rope import defer_fused_qkv_rope + + deferred = defer_fused_qkv_rope( + qkv, + q_freqs, + k_freqs, + qkv_split_arg_list, + start_positions, + tensor_format, + interleaved, + cp_size, + ) + if deferred is not None: + return deferred + return FusedQKVRoPEFunc.apply( qkv, q_freqs, diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 7eede479fb9..627e7f89b56 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -5,6 +5,7 @@ """Tensor class with FP8 data""" from __future__ import annotations from typing import Any, Dict, Optional, Tuple, Iterable, Union +import os import warnings import torch from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState @@ -195,6 +196,14 @@ def supports_only_rowwise_all_gather(self) -> bool: return True +_TRITON_CURRENT_SCALING = os.environ.get("NVTE_TRITON_CURRENT_SCALING_QUANTIZE", "1") == "1" +_TRITON_CURRENT_SCALING_MIN_NUMEL = 1 << 22 +_FP8_TORCH_DTYPES = { + DType.kFloat8E4M3: torch.float8_e4m3fn, + DType.kFloat8E5M2: torch.float8_e5m2, +} + + class Float8CurrentScalingQuantizer(Quantizer): """Builder class for FP8 tensors with per-tensor current scaling @@ -312,8 +321,40 @@ def update_quantized( return dst + def _triton_quantize(self, tensor: torch.Tensor) -> Optional[QuantizedTensor]: + """Two-phase Triton amax + cast for large rowwise-only tensors; None if unsupported.""" + if ( + not _TRITON_CURRENT_SCALING + or tensor.dtype != torch.bfloat16 + or not tensor.is_cuda + or not tensor.is_contiguous() + or tensor.numel() < _TRITON_CURRENT_SCALING_MIN_NUMEL + or self.with_amax_reduction + or self.columnwise_usage + or not self.rowwise_usage + or self.dtype not in _FP8_TORCH_DTYPES + ): + return None + out = self.make_empty(tensor.shape, dtype=tensor.dtype, device=tensor.device) + if out._data is None or out._transpose is not None or out._scale_inv.numel() != 1: + return None + # pylint: disable=import-outside-toplevel + from ..triton.current_scaling import current_scaling_quantize + + current_scaling_quantize( + tensor, + out._data.view(_FP8_TORCH_DTYPES[self.dtype]), + out._scale_inv, + eps=float(self.amax_epsilon), + force_pow_2_scales=self.force_pow_2_scales, + ) + return out + def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" + out = self._triton_quantize(tensor) + if out is not None: + return out return tex.quantize(tensor, self) def calibrate(self, tensor: torch.Tensor) -> None: diff --git a/transformer_engine/pytorch/triton/current_scaling.py b/transformer_engine/pytorch/triton/current_scaling.py new file mode 100644 index 00000000000..82aeb61ae0c --- /dev/null +++ b/transformer_engine/pytorch/triton/current_scaling.py @@ -0,0 +1,61 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrapper for the two-phase FP8 current-scaling quantization Triton kernel.""" + +from typing import Dict, Tuple + +import torch +import triton + +from transformer_engine.common.triton.current_scaling import current_scaling_quantize_kernel + +_BLOCK = 8192 +_NUM_WARPS = 8 +_CTAS_PER_SM = 4 +_FP8_MAX = {torch.float8_e4m3fn: 448.0, torch.float8_e5m2: 57344.0} + +# One [amax, counter, scale, pad] workspace per (device, stream): the last CTA of every launch +# re-zeroes it, but launches on different streams must not share one. +_WORKSPACES: Dict[Tuple[int, int], torch.Tensor] = {} +_NUM_SMS: Dict[int, int] = {} + + +def _workspace(device: torch.device) -> torch.Tensor: + key = (device.index, torch.cuda.current_stream(device).cuda_stream) + ws = _WORKSPACES.get(key) + if ws is None: + ws = _WORKSPACES[key] = torch.zeros(4, dtype=torch.float32, device=device) + return ws + + +def current_scaling_quantize( + inp: torch.Tensor, + out: torch.Tensor, + scale_inv: torch.Tensor, + *, + eps: float, + force_pow_2_scales: bool, +) -> None: + """Per-tensor current-scaling FP8 quantization of a contiguous tensor. + + ``out`` is a float8_e4m3fn/e5m2 tensor of the same number of elements and ``scale_inv`` a + 1-element fp32 tensor; both are written. Bitwise equal to TE's amax + compute-scale + cast. + """ + n = inp.numel() + device = inp.device + if device.index not in _NUM_SMS: + _NUM_SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count + ws = _workspace(device) + n_blk = triton.cdiv(n, _BLOCK) + common = dict( + BLOCK=_BLOCK, + FP8MAX=_FP8_MAX[out.dtype], + POW2=bool(force_pow_2_scales), + num_warps=_NUM_WARPS, + ) + current_scaling_quantize_kernel[(min(n_blk, _NUM_SMS[device.index] * _CTAS_PER_SM),)]( + inp, out, ws, scale_inv, n, eps, CAST=False, **common + ) + current_scaling_quantize_kernel[(n_blk,)](inp, out, ws, scale_inv, n, eps, CAST=True, **common) diff --git a/transformer_engine/pytorch/triton/rope_fp8.py b/transformer_engine/pytorch/triton/rope_fp8.py new file mode 100644 index 00000000000..4b72c6f4060 --- /dev/null +++ b/transformer_engine/pytorch/triton/rope_fp8.py @@ -0,0 +1,99 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrappers for the fused RoPE + FP8 current-scaling Q/K/V quantization kernels.""" + +from typing import Dict, Tuple + +import torch +import triton + +from transformer_engine.common.triton.rope_fp8 import ( + rope_qkv_quantize_kernel, + rope_trig_table_kernel, +) +from .current_scaling import _FP8_MAX, _CTAS_PER_SM, _NUM_SMS, _workspace + +_AMAX_WARPS = 2 +_CAST_WARPS = 4 +_AMAX_STAGES = 3 + +# (cos, sin) per device, valid for one freqs tensor object at one version. The entry holds the +# tensor, so its address cannot be reused while the entry is alive. +_TRIG_CACHE: Dict[int, Tuple[torch.Tensor, int, torch.Tensor, torch.Tensor]] = {} + + +def rope_trig_tables(freqs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """cos/sin (fp32, freqs' shape) of contiguous fp32 RoPE angles, bitwise equal to TE's + sincosf. All layers of a step share one rotary embedding tensor, so this runs once per step.""" + ent = _TRIG_CACHE.get(freqs.device.index) + if ent is not None and ent[0] is freqs and ent[1] == freqs._version: + return ent[2], ent[3] + cos = torch.empty_like(freqs) + sin = torch.empty_like(freqs) + n = freqs.numel() + rope_trig_table_kernel[(triton.cdiv(n, 1024),)](freqs, cos, sin, n, BLOCK=1024, num_warps=4) + _TRIG_CACHE[freqs.device.index] = (freqs, freqs._version, cos, sin) + return cos, sin + + +def rope_qkv_quantize( + mixed_qkv: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + out: torch.Tensor, + scale_inv: torch.Tensor, + num_q_per_group: int, + *, + eps: float, + force_pow_2_scales: bool, +) -> None: + """Apply RoPE to Q/K of ``mixed_qkv`` [s, b, g, (hq + 2) * d] (bf16) and quantize Q, K, V + with one current-scaling scale into ``out`` (flat float8, packed [Q | K | V] with Q as + (s, b, g * hq, d) and K, V as (s, b, g, d)). ``scale_inv`` is a 1-element fp32 tensor.""" + s, b, g, width = mixed_qkv.shape + d = cos.shape[-1] + n_tok = s * b + assert width == (num_q_per_group + 2) * d, "mixed_qkv last dim must be (hq + 2) * d" + assert out.numel() == n_tok * g * width, "out must hold the packed [Q | K | V] payload" + assert mixed_qkv.numel() < 2**31, "tensor too large for 32-bit indexing" + device = mixed_qkv.device + if device.index not in _NUM_SMS: + _NUM_SMS[device.index] = torch.cuda.get_device_properties(device).multi_processor_count + ws = _workspace(device) + common = dict( + B=b, + G=g, + HQ=num_q_per_group, + D=d, + FP8MAX=_FP8_MAX[out.dtype], + POW2=bool(force_pow_2_scales), + ) + rope_qkv_quantize_kernel[(min(n_tok, _NUM_SMS[device.index] * _CTAS_PER_SM),)]( + mixed_qkv, + cos, + sin, + out, + ws, + scale_inv, + n_tok, + eps, + CAST=False, + NUM_STAGES=_AMAX_STAGES, + num_warps=_AMAX_WARPS, + **common, + ) + rope_qkv_quantize_kernel[(n_tok,)]( + mixed_qkv, + cos, + sin, + out, + ws, + scale_inv, + n_tok, + eps, + CAST=True, + num_warps=_CAST_WARPS, + **common, + )