Fused QKV ROPE - CS FP8 quantization

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

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

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)

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];
+    #                                  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)