From 703a13bbc7fd3f8dbd79cc50ef54fafff2dae702 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Mon, 9 Nov 2026 14:47:02 -0800 Subject: [PATCH 1/2] [Common][PyTorch] MLA q/kv RoPE + MXFP8 in attention's scale layout, and MXFP8 RoPE backward Multi-latent attention with MXFP8 attention runs, per layer and microbatch, the query RoPE, its MXFP8 quantization, the key/value RoPE + MXFP8 kernel and then mxfp8_transpose_swizzle, which permutes the q/k/v scaling factors to BHSD and swizzles them for the GEMM. In the backward, the query RoPE backward and the key/value split backward write BF16 gradients of the up-projection outputs, which the up-projections then quantize to MXFP8. Add kernels that produce the same bytes directly: - nvte_fused_mla_q_rope_mxfp8_attn: query RoPE + MXFP8, scales in attention's layout (BHSD-permuted, GEMM-swizzled); - nvte_fused_mla_kv_rope_mxfp8_attn: the key/value split, key RoPE and MXFP8 of nvte_fused_mla_kv_rope_mxfp8, with the scales in attention's layout; - nvte_fused_mla_q_rope_backward: the query RoPE backward in place, with Megatron-LM's BF16 roundings; - nvte_fused_mla_q_rope_backward_mxfp8 and nvte_fused_mla_kv_split_backward_mxfp8: the query RoPE backward and the key/value split backward written as MXFP8 (row-wise and column-wise, GEMM-swizzled scales), the grad_output the up-projections would have quantized. PyTorch: fused_mla_q_rope_mxfp8, fused_mla_kv_rope_mxfp8(..., transpose_swizzle_scales=True), fused_mla_q_rope_backward_, fused_mla_q_rope_backward_mxfp8 and fused_mla_kv_split_backward_mxfp8 in attention/dot_product_attention/utils.py. Signed-off-by: Siddhartha Raman S Co-Authored-By: Claude Opus 5.5 --- tests/pytorch/attention/test_fused_mla.py | 266 ++++- transformer_engine/common/CMakeLists.txt | 1 + .../fused_rope/fused_mla_rope_mxfp8_attn.cu | 986 ++++++++++++++++++ .../include/transformer_engine/fused_rope.h | 105 ++ .../attention/dot_product_attention/utils.py | 175 +++- transformer_engine/pytorch/csrc/extensions.h | 16 + .../pytorch/csrc/extensions/apply_rope.cpp | 174 ++++ .../pytorch/csrc/extensions/pybind.cpp | 17 + 8 files changed, 1736 insertions(+), 4 deletions(-) create mode 100644 transformer_engine/common/fused_rope/fused_mla_rope_mxfp8_attn.cu diff --git a/tests/pytorch/attention/test_fused_mla.py b/tests/pytorch/attention/test_fused_mla.py index 1fcd9a8e218..bc10a0bf0a4 100644 --- a/tests/pytorch/attention/test_fused_mla.py +++ b/tests/pytorch/attention/test_fused_mla.py @@ -8,6 +8,13 @@ - fused_mla_kv_rope_mxfp8: key/value split, key RoPE and MXFP8 quantization. The reference applies Megatron-LM's MLA key/value RoPE (the Triton kernel in mla_rope_utils) and quantizes key and value with mxfp8_quantize_only; the fused kernel must produce the same bytes. +- fused_mla_q_rope_mxfp8 and fused_mla_kv_rope_mxfp8(transpose_swizzle_scales=True): query RoPE + / key-value split and RoPE with MXFP8 quantization in the attention scale layout. Same bytes as + the RoPE kernels, mxfp8_quantize_only and mxfp8_transpose_swizzle. +- fused_mla_q_rope_backward_, fused_mla_q_rope_backward_mxfp8 and + fused_mla_kv_split_backward_mxfp8: same bytes as Megatron-LM's MLA RoPE backward kernels, + followed (for the MXFP8 ones) by the MXFP8 quantizer with optimize_for_gemm, as TE's linear + layers quantize their grad_output. Run: pytest tests/pytorch/attention/test_fused_mla.py -v @@ -25,14 +32,19 @@ from transformer_engine.pytorch.attention import FusedMLAQUpProjFunction, FusedMLAQUpProjRopeQuant from transformer_engine.pytorch.attention.dot_product_attention.utils import ( fused_mla_kv_rope_mxfp8, + fused_mla_kv_split_backward_mxfp8, + fused_mla_q_rope_backward_, + fused_mla_q_rope_backward_mxfp8, + fused_mla_q_rope_mxfp8, mxfp8_quantize_only, + mxfp8_transpose_swizzle, ) from transformer_engine.pytorch.cpp_extensions import general_gemm as _fused_general_gemm from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor _current_file = pathlib.Path(__file__).resolve() sys.path = [str(_current_file.parent)] + sys.path -from mla_rope_utils import HAVE_TRITON, apply_mla_rope_kv, build_rope_tables +from mla_rope_utils import HAVE_TRITON, apply_mla_rope_kv, apply_mla_rope_q, build_rope_tables # All tests in this file require GB200/SM100 hardware. pytestmark = pytest.mark.launch_on_gb200 @@ -462,3 +474,255 @@ def test_inputs_on_non_current_device(): _assert_identical(value, ref_value, "value") with pytest.raises(RuntimeError, match="same device"): fused_mla_kv_rope_mxfp8(kv, k_rope.to("cuda:0"), cos, sin, HEAD_DIM_V) + + +# --- MXFP8 attention scale layout and the RoPE backward kernels ----------------------------------- + + +def _rope_tables(s, dtype=torch.bfloat16, rows=None): + cos, sin = build_rope_tables(rows or s, emb_dim=HEAD_DIM_ROPE, device="cuda") + return (cos * 1.2).to(dtype), (sin * 1.2).to(dtype) + + +def _q_input(s, b, h, seed=4321): + torch.manual_seed(seed) + q = (torch.randn(s, b, h, HEAD_DIM, device="cuda") * 4).to(torch.bfloat16) + q[3, 0, 0, :32] = 0 + q[11, b - 1, h - 1, HEAD_DIM_NOPE + 9] = 2.0e4 + q[17, 0, 1, 70] = -3.0e4 + return q + + +def _q_reference(q, cos, sin): + query = apply_mla_rope_q(q.clone(), cos_table=cos, sin_table=sin) + (query,) = mxfp8_quantize_only([(query.contiguous(), _kv_quantizer())], "sbhd") + mxfp8_transpose_swizzle([query], "sbhd") + return query + + +def _assert_swizzled_identical(actual, expected, name): + """Same bytes, with the scale_invs compared as flat buffers.""" + assert actual._with_gemm_swizzled_scales and expected._with_gemm_swizzled_scales, name + for buffer in MXFP8_BUFFERS: + a, e = getattr(actual, buffer), getattr(expected, buffer) + assert a.numel() == e.numel(), f"{name}.{buffer}: {a.numel()} vs {e.numel()} bytes" + a, e = a.reshape(-1).view(torch.uint8), e.reshape(-1).view(torch.uint8) + mismatches = int((a != e).sum()) + assert mismatches == 0, f"{name}.{buffer}: {mismatches} of {e.numel()} bytes differ" + + +ATTN_SHAPES = [(4096, 1, 128), (2048, 2, 64), (512, 1, 16), (256, 3, 6), (128, 2, 2)] + + +@requires_kv_rope_kernel +@pytest.mark.parametrize("s, b, h", ATTN_SHAPES) +def test_q_rope_mxfp8_matches_rope_quantize_swizzle(s, b, h): + q = _q_input(s, b, h) + cos, sin = _rope_tables(s) + expected = _q_reference(q, cos, sin) + q_before = q.clone() + query = fused_mla_q_rope_mxfp8(q, cos.view(s, 1, 1, -1), sin.view(s, 1, 1, -1)) + assert query.shape == q.shape + assert torch.equal(q, q_before), "the input must not be modified" + _assert_swizzled_identical(query, expected, "query") + assert tuple(query._rowwise_scale_inv.shape) == tuple(expected._rowwise_scale_inv.shape) + assert tuple(query._columnwise_scale_inv.shape) == tuple(expected._columnwise_scale_inv.shape) + + +@requires_kv_rope_kernel +@pytest.mark.parametrize("s, b, h", ATTN_SHAPES) +def test_kv_rope_mxfp8_transpose_swizzle_scales(s, b, h): + kv, k_rope, cos, sin = _kv_rope_inputs(s, b, h, torch.bfloat16) + ref_key, ref_value = _kv_rope_reference(kv, k_rope, cos, sin) + mxfp8_transpose_swizzle([ref_key, ref_value], "sbhd") + key, value = fused_mla_kv_rope_mxfp8( + kv, k_rope, cos, sin, HEAD_DIM_V, transpose_swizzle_scales=True + ) + _assert_swizzled_identical(key, ref_key, "key") + _assert_swizzled_identical(value, ref_value, "value") + # The compact-scale kernel followed by the scale permute and swizzle gives the same bytes. + key2, value2 = fused_mla_kv_rope_mxfp8(kv, k_rope, cos, sin, HEAD_DIM_V) + mxfp8_transpose_swizzle([key2, value2], "sbhd") + _assert_swizzled_identical(key, key2, "key") + _assert_swizzled_identical(value, value2, "value") + + +@requires_kv_rope_kernel +def test_attn_layout_non_finite_and_long_tables(): + """Inf / NaN inputs, a contiguous k_rope and cos / sin with more rows than tokens.""" + s, b, h = 256, 2, 4 + q = _q_input(s, b, h) + q[5, 0, 0, 40] = float("inf") + q[70, 1, 3, HEAD_DIM_NOPE + 2] = float("nan") + kv, k_rope, _, _ = _kv_rope_inputs(s, b, h, torch.bfloat16) + kv[3, 0, 0, 40] = float("-inf") + kv[96:128, 1, 3, HEAD_DIM_NOPE + 32 : HEAD_DIM_NOPE + 64] = float("nan") + k_rope = k_rope.squeeze(-2).contiguous() + cos, sin = _rope_tables(s, rows=2 * s) + _assert_swizzled_identical( + fused_mla_q_rope_mxfp8(q, cos, sin), _q_reference(q, cos[:s], sin[:s]), "query" + ) + ref_key, ref_value = _kv_rope_reference(kv, k_rope.unsqueeze(-2), cos[:s], sin[:s]) + mxfp8_transpose_swizzle([ref_key, ref_value], "sbhd") + key, value = fused_mla_kv_rope_mxfp8( + kv, k_rope, cos, sin, HEAD_DIM_V, transpose_swizzle_scales=True + ) + _assert_swizzled_identical(key, ref_key, "key") + _assert_swizzled_identical(value, ref_value, "value") + + +@requires_kv_rope_kernel +@pytest.mark.parametrize( + "s, h, cos_dtype, message", + [ + (96, 4, torch.bfloat16, "multiple of 128"), + (128, 3, torch.bfloat16, "must be even"), + (128, 4, torch.float32, "must be BF16"), + ], + ids=["s", "h", "fp32_cos"], +) +def test_attn_layout_unsupported_inputs(s, h, cos_dtype, message): + q = _q_input(s, 1, h) + kv, k_rope, _, _ = _kv_rope_inputs(s, 1, h, torch.bfloat16) + cos, sin = _rope_tables(s, cos_dtype) + with pytest.raises((RuntimeError, ValueError), match=message): + fused_mla_q_rope_mxfp8(q, cos, sin) + with pytest.raises((RuntimeError, ValueError), match=message): + fused_mla_kv_rope_mxfp8(kv, k_rope, cos, sin, HEAD_DIM_V, transpose_swizzle_scales=True) + + +def _grad_quantizer(): + quantizer = _kv_quantizer() + quantizer.optimize_for_gemm = True # As TE's linear layers quantize grad_output. + return quantizer + + +def _q_rope_backward_reference(dq, cos, sin): + s, b, h, _ = dq.shape + q = torch.zeros(s, b, h, HEAD_DIM, device="cuda", dtype=torch.bfloat16, requires_grad=True) + out = apply_mla_rope_q(q.clone(), cos_table=cos, sin_table=sin) + out.backward(dq.clone()) + return q.grad + + +def _kv_split_backward_reference(dk, dv, cos, sin): + s, b, h, _ = dk.shape + kv = torch.zeros( + s, b, h, HEAD_DIM_NOPE + HEAD_DIM_V, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + k_rope = torch.zeros(s, b, 1, HEAD_DIM_ROPE, device="cuda", dtype=torch.bfloat16) + k_rope.requires_grad_() + key, value = apply_mla_rope_kv(kv, k_rope, cos_table=cos, sin_table=sin) + torch.autograd.backward([key, value], [dk, dv]) + return kv.grad + + +def _grad_inputs(s, b, h, seed=99): + torch.manual_seed(seed) + dq = (torch.randn(s, b, h, HEAD_DIM, device="cuda") * 1e-3).to(torch.bfloat16) + dk = (torch.randn(s, b, h, HEAD_DIM, device="cuda") * 1e-3).to(torch.bfloat16) + dv = (torch.randn(s, b, h, HEAD_DIM_V, device="cuda") * 1e-3).to(torch.bfloat16) + for t in (dq, dk, dv): + t[1, 0, 0, :32] = 0 + t[33, b - 1, h - 1, 5] = 30.0 + t[min(64, s - 1), 0, 1, 100] = -1.0e-30 + return dq, dk, dv + + +@requires_kv_rope_kernel +@pytest.mark.parametrize("s, b, h", ATTN_SHAPES + [(64, 2, 4)]) +def test_q_rope_backward_in_place(s, b, h): + dq, _, _ = _grad_inputs(s, b, h) + cos, sin = _rope_tables(s) + expected = _q_rope_backward_reference(dq, cos, sin) + out = fused_mla_q_rope_backward_(dq, cos, sin) + assert out.data_ptr() == dq.data_ptr() + assert torch.equal(dq.view(torch.int16), expected.view(torch.int16)) + + +@requires_kv_rope_kernel +@pytest.mark.parametrize("s, b, h", ATTN_SHAPES + [(64, 2, 4)]) +def test_q_rope_backward_mxfp8(s, b, h): + dq, _, _ = _grad_inputs(s, b, h) + cos, sin = _rope_tables(s) + expected_bf16 = _q_rope_backward_reference(dq, cos, sin) + expected = _grad_quantizer()(expected_bf16.view(s * b, h * HEAD_DIM)) + dq_before = dq.clone() + grad = fused_mla_q_rope_backward_mxfp8(dq, cos, sin) + assert torch.equal(dq, dq_before), "dq must not be modified" + assert grad.shape == (s, b, h * HEAD_DIM) + _assert_swizzled_identical(grad, expected, "grad") + + +@requires_kv_rope_kernel +@pytest.mark.parametrize("s, b, h", ATTN_SHAPES + [(64, 2, 4)]) +def test_kv_split_backward_mxfp8(s, b, h): + _, dk, dv = _grad_inputs(s, b, h) + cos, sin = _rope_tables(s) + expected_bf16 = _kv_split_backward_reference(dk, dv, cos, sin) + expected = _grad_quantizer()(expected_bf16.view(s * b, -1)) + grad = fused_mla_kv_split_backward_mxfp8(dk, dv, HEAD_DIM_NOPE) + assert grad.shape == (s, b, h * (HEAD_DIM_NOPE + HEAD_DIM_V)) + _assert_swizzled_identical(grad, expected, "grad") + + +@requires_kv_rope_kernel +def test_backward_unsupported_shapes(): + _, dk, dv = _grad_inputs(96, 1, 4) + cos, sin = _rope_tables(96) + with pytest.raises(RuntimeError, match="multiple of 128"): + fused_mla_q_rope_backward_mxfp8(dk, cos, sin) + with pytest.raises(RuntimeError, match="multiple of 128"): + fused_mla_kv_split_backward_mxfp8(dk, dv, HEAD_DIM_NOPE) + _, dk, dv = _grad_inputs(128, 1, 3) + with pytest.raises(RuntimeError, match="must be even"): + fused_mla_kv_split_backward_mxfp8(dk, dv, HEAD_DIM_NOPE) + + +@requires_kv_rope_kernel +def test_cuda_graph_replay(): + """All five kernels captured in a CUDA graph and replayed on new inputs.""" + s, b, h = 256, 2, 8 + cos, sin = _rope_tables(s) + q = _q_input(s, b, h) + kv, k_rope, _, _ = _kv_rope_inputs(s, b, h, torch.bfloat16) + dq, dk, dv = _grad_inputs(s, b, h) + dq_inplace = dq.clone() + + def run(): + query = fused_mla_q_rope_mxfp8(q, cos, sin) + key, value = fused_mla_kv_rope_mxfp8( + kv, k_rope, cos, sin, HEAD_DIM_V, transpose_swizzle_scales=True + ) + fused_mla_q_rope_backward_(dq_inplace, cos, sin) + dq_grad = fused_mla_q_rope_backward_mxfp8(dq, cos, sin) + dkv_grad = fused_mla_kv_split_backward_mxfp8(dk, dv, HEAD_DIM_NOPE) + return query, key, value, dq_grad, dkv_grad + + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() # Warm-up outside capture. + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = run() + + for seed in (7, 8): + q.copy_(_q_input(s, b, h, seed=seed)) + new_kv, new_k_rope, _, _ = _kv_rope_inputs(s, b, h, torch.bfloat16, seed=seed) + kv.copy_(new_kv) + k_rope.copy_(new_k_rope) + new_dq, new_dk, new_dv = _grad_inputs(s, b, h, seed=seed) + dq.copy_(new_dq) + dq_inplace.copy_(new_dq) + dk.copy_(new_dk) + dv.copy_(new_dv) + graph.replay() + torch.cuda.synchronize() + dq_inplace_graph = dq_inplace.clone() + dq_inplace.copy_(new_dq) + eager = run() + for name, a, e in zip(("query", "key", "value", "dq_grad", "dkv_grad"), outputs, eager): + _assert_swizzled_identical(a, e, name) + assert torch.equal(dq_inplace_graph, dq_inplace) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f90f74449c5..b6357404062 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -365,6 +365,7 @@ list(APPEND transformer_engine_cuda_arch_specific_sources cast/cast_grouped_dbias.cu cast/mxfp8/requantize_mxfp8.cu fused_rope/fused_mla_kv_rope_mxfp8.cu + fused_rope/fused_mla_rope_mxfp8_attn.cu gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu hadamard_transform/graph_safe_group_hadamard_transform.cu diff --git a/transformer_engine/common/fused_rope/fused_mla_rope_mxfp8_attn.cu b/transformer_engine/common/fused_rope/fused_mla_rope_mxfp8_attn.cu new file mode 100644 index 00000000000..1b49fbf66ab --- /dev/null +++ b/transformer_engine/common/fused_rope/fused_mla_rope_mxfp8_attn.cu @@ -0,0 +1,986 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file fused_mla_rope_mxfp8_attn.cu + * \brief Multi-latent attention RoPE fused with MXFP8 quantization, writing the scaling factor + * layouts that MXFP8 attention and the up-projection GEMMs consume. + * + * Forward, for MXFP8 attention (DeepSeek-V3 head dimensions: 128 + 64 query/key, 128 value): + * - query: RoPE of the query's last 64 dimensions and MXFP8 quantization; + * - key/value: split of the key/value up-projection output, RoPE of the shared key RoPE + * dimensions and MXFP8 quantization of key and value. + * Both write the data in SBHD and the scaling factors permuted to BHSD and swizzled for the GEMM, + * the bytes of mxfp8_quantize_only followed by mxfp8_transpose_swizzle (or + * fused_mla_kv_rope_mxfp8 followed by mxfp8_transpose_swizzle), so the scale permute and swizzle + * are not separate kernels. + * + * Backward: + * - the query RoPE backward in place (BF16, Megatron-LM's roundings); + * - the query RoPE backward and the key/value split backward fused with the MXFP8 quantization + * of the result, which is the up-projections' grad_output: row-wise and column-wise data and + * GEMM-swizzled scaling factors, without writing the BF16 gradient. + * + * Quantization is that of the MXFP8 quantizer (E4M3): amax of 32 BF16 values, E8M0 scale + * rounded up from amax / 448, data = x * 2^-scale in BF16 converted to E4M3 with saturation. + * + * Work unit of the quantizing kernels: 32 tokens x 64 columns, one warp. Lane 8 * r + k holds + * the 16-byte chunk k (8 columns) of tokens 4 * i + r, i < 8. Row-wise amax: in-lane over 8 + * values and 2 shuffles; column-wise: in-lane over 8 tokens and 2 shuffles. A 128-token tile's + * scaling factors are whole 512-byte blocks of the swizzled layout; they are assembled in shared + * memory and written with 16-byte stores. + */ + +#include +#include +#include + +#include +#include +#include +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" +#include "../util/logging.h" +#include "../util/ptx.cuh" +#include "../util/ptx_arch_spec.cuh" +#include "../utils.cuh" + +namespace transformer_engine { +namespace { + +constexpr int kNopeDim = 128; // Query/key dimensions without RoPE. +constexpr int kRopeDim = 64; // RoPE dimensions of the query/key. +constexpr int kQKDim = kNopeDim + kRopeDim; // 192 +constexpr int kValueDim = 128; +constexpr int kKVDim = kNopeDim + kValueDim; // 256, per head of the key/value up-projection +constexpr int kTileTokens = 128; // Tokens of a swizzled scaling factor block. + +// --------------------------------------------------------------------------------------------- +// Quantization of one unit (32 tokens x 64 columns) held by a warp. +// --------------------------------------------------------------------------------------------- + +__device__ __forceinline__ uint32_t habs2(uint32_t w) { + __nv_bfloat162 v = __habs2(*reinterpret_cast<__nv_bfloat162 *>(&w)); + return *reinterpret_cast(&v); +} + +__device__ __forceinline__ uint32_t hmax2(uint32_t a, uint32_t b) { + __nv_bfloat162 v = + __hmax2(*reinterpret_cast<__nv_bfloat162 *>(&a), *reinterpret_cast<__nv_bfloat162 *>(&b)); + return *reinterpret_cast(&v); +} + +__device__ __forceinline__ float bf16_lo(uint32_t w) { return __uint_as_float(w << 16); } +__device__ __forceinline__ float bf16_hi(uint32_t w) { return __uint_as_float(w & 0xFFFF0000u); } + +// E8M0 scales of two FP32 values (rounded up, saturating): byte 0 <- lo, byte 1 <- hi. +__device__ __forceinline__ uint32_t e8m0x2(const float lo, const float hi) { + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + uint16_t out; + asm volatile("{\n cvt.rp.satfinite.ue8m0x2.f32 %0, %1, %2;\n}" : "=h"(out) : "f"(hi), "f"(lo)); + return out; + } else { + return ptx::float_to_e8m0(lo) | (static_cast(ptx::float_to_e8m0(hi)) << 8); + } +} + +// BF16 2^-e of an E8M0 exponent, as ptx::exp2f_rcp. +__device__ __forceinline__ uint32_t exp2f_rcp_bf16_bits(const uint32_t e) { + if (e == 255) return 0x7fffu; + if (e == 254) return 0x0040u; + return ((254 - e) << 7) & 0xFFFFu; +} + +// BF16 reciprocal scales of the two E8M0 exponents in bytes 0 and 1 of e2, as a bf16x2 +// (low half <- byte 0). +__device__ __forceinline__ uint32_t rcp_scales2(const uint32_t e2) { + const uint32_t spread = __byte_perm(e2, 0, 0x7170); // byte 0 -> bits 0..7, byte 1 -> 16..23 + uint32_t bits = 0x7F007F00u - spread * 0x80u; + if (__builtin_expect((e2 & 0xFEu) == 0xFEu || (e2 & 0xFE00u) == 0xFE00u, 0)) { + bits = exp2f_rcp_bf16_bits(e2 & 0xFF) | (exp2f_rcp_bf16_bits((e2 >> 8) & 0xFF) << 16); + } + return bits; +} + +// Four BF16 values (two words) times one bf16x2 scale, converted to E4M3. +__device__ __forceinline__ uint32_t cvt_4x(const uint32_t x01, const uint32_t x23, + const uint32_t scale) { + const uint2 in = make_uint2(x01, x23); + uint32_t out; + ptx::mul_cvt_4x(reinterpret_cast(out), + reinterpret_cast(in), + reinterpret_cast(scale)); + return out; +} + +// Two BF16 pairs, each times its own bf16x2 scale, converted to E4M3. +__device__ __forceinline__ uint32_t cvt_2x2(const uint32_t in0, const uint32_t s0, + const uint32_t in1, const uint32_t s1) { + uint32_t out; + ptx::mul_cvt_4x( + reinterpret_cast(out), reinterpret_cast(in0), + reinterpret_cast(s0), reinterpret_cast(in1), + reinterpret_cast(s1)); + return out; +} + +// A quantized unit, for lane 8 * r + k: row-wise and column-wise E4M3 of tokens 4 * i + r, +// columns 8 * k .. 8 * k + 7. +struct QuantizedUnit { + uint2 row[8], col[8]; + uint32_t rs[8]; // Lanes with k == 0: row-wise scales (blocks 0, 1 of the unit) of token 4i + r. + uint2 cs; // Column-wise scales of columns 8 * k .. 8 * k + 7. +}; + +__device__ __forceinline__ void quantize(const uint32_t (&v)[8][4], QuantizedUnit &o) { + constexpr float kMaxNormRcp = Quantized_Limits::max_norm_rcp; + uint32_t a[8][4]; +#pragma unroll + for (int i = 0; i < 8; ++i) { +#pragma unroll + for (int j = 0; j < 4; ++j) a[i][j] = habs2(v[i][j]); + } + // Row-wise: amax per token over the lane's 8 values (two tokens per word), then over the + // block's 4 lanes. +#pragma unroll + for (int i = 0; i < 8; i += 2) { + const uint32_t m0 = hmax2(hmax2(a[i][0], a[i][1]), hmax2(a[i][2], a[i][3])); + const uint32_t m1 = hmax2(hmax2(a[i + 1][0], a[i + 1][1]), hmax2(a[i + 1][2], a[i + 1][3])); + // Low half: token i, high half: token i + 1. + uint32_t m = hmax2(__byte_perm(m0, m1, 0x5410), __byte_perm(m0, m1, 0x7632)); + m = hmax2(m, __shfl_xor_sync(0xffffffffu, m, 1)); + m = hmax2(m, __shfl_xor_sync(0xffffffffu, m, 2)); + const uint32_t e2 = e8m0x2(bf16_lo(m) * kMaxNormRcp, bf16_hi(m) * kMaxNormRcp); + const uint32_t sc2 = rcp_scales2(e2); + const uint32_t s0 = __byte_perm(sc2, 0, 0x1010), s1 = __byte_perm(sc2, 0, 0x3232); + o.row[i] = make_uint2(cvt_4x(v[i][0], v[i][1], s0), cvt_4x(v[i][2], v[i][3], s0)); + o.row[i + 1] = + make_uint2(cvt_4x(v[i + 1][0], v[i + 1][1], s1), cvt_4x(v[i + 1][2], v[i + 1][3], s1)); + const uint32_t e2_block1 = __shfl_xor_sync(0xffffffffu, e2, 4); // Lane k + 4. + o.rs[i] = __byte_perm(e2, e2_block1, 0x0040); + o.rs[i + 1] = __byte_perm(e2, e2_block1, 0x0051); + } + // Column-wise. + uint32_t s[4], e[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + uint32_t m = hmax2(hmax2(hmax2(a[0][j], a[1][j]), hmax2(a[2][j], a[3][j])), + hmax2(hmax2(a[4][j], a[5][j]), hmax2(a[6][j], a[7][j]))); + m = hmax2(m, __shfl_xor_sync(0xffffffffu, m, 8)); + m = hmax2(m, __shfl_xor_sync(0xffffffffu, m, 16)); + e[j] = e8m0x2(bf16_lo(m) * kMaxNormRcp, bf16_hi(m) * kMaxNormRcp); + s[j] = rcp_scales2(e[j]); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + o.col[i] = + make_uint2(cvt_2x2(v[i][0], s[0], v[i][1], s[1]), cvt_2x2(v[i][2], s[2], v[i][3], s[3])); + } + o.cs = make_uint2(e[0] | (e[1] << 16), e[2] | (e[3] << 16)); +} + +// Megatron-LM's MLA RoPE with BF16 cos / sin, 8 rotation pairs (x1, x2) -> 8 outputs (4 words) +// of the left or the right half: left = fma(x1, cos_l, -(x2 * sin_l)), +// right = fma(x1, sin_r, x2 * cos_r), each operation rounded to BF16. +__device__ __forceinline__ void rope8(const uint4 x_lo, const uint4 x_hi, const uint4 c, + const uint4 sn, const bool right, uint32_t (&w)[4]) { + const uint32_t xs[8] = {x_lo.x, x_lo.y, x_lo.z, x_lo.w, x_hi.x, x_hi.y, x_hi.z, x_hi.w}; + const uint32_t cw[4] = {c.x, c.y, c.z, c.w}, sw[4] = {sn.x, sn.y, sn.z, sn.w}; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t x1u = __byte_perm(xs[2 * j], xs[2 * j + 1], 0x5410); + const uint32_t x2u = __byte_perm(xs[2 * j], xs[2 * j + 1], 0x7632); + const __nv_bfloat162 x1 = *reinterpret_cast(&x1u); + const __nv_bfloat162 x2 = *reinterpret_cast(&x2u); + const __nv_bfloat162 cj = *reinterpret_cast(&cw[j]); + const __nv_bfloat162 sj = *reinterpret_cast(&sw[j]); + const __nv_bfloat162 y = + right ? __hfma2(x1, sj, __hmul2(x2, cj)) : __hfma2(x1, cj, __hneg2(__hmul2(x2, sj))); + w[j] = *reinterpret_cast(&y); + } +} + +// Megatron-LM's MLA query RoPE backward for 2 x 2 rotation pairs: left = dq[nope : nope + 32], +// right = dq[nope + 32 : nope + 64]; out[2p] = fma(left_p, cos_l_p, right_p * sin_r_p), +// out[2p + 1] = fma(right_p, cos_r_p, -(left_p * sin_l_p)), each operation rounded to BF16. +// Writes the four output words (interleaved pairs). +__device__ __forceinline__ void rope_bwd4(const uint2 left, const uint2 right, const uint2 cl, + const uint2 sl, const uint2 cr, const uint2 sr, + uint32_t (&o)[4]) { + const uint32_t lw[2] = {left.x, left.y}, rw[2] = {right.x, right.y}; + const uint32_t clw[2] = {cl.x, cl.y}, slw[2] = {sl.x, sl.y}; + const uint32_t crw[2] = {cr.x, cr.y}, srw[2] = {sr.x, sr.y}; +#pragma unroll + for (int w = 0; w < 2; ++w) { + const __nv_bfloat162 xl = *reinterpret_cast(&lw[w]); + const __nv_bfloat162 xr = *reinterpret_cast(&rw[w]); + const __nv_bfloat162 c_l = *reinterpret_cast(&clw[w]); + const __nv_bfloat162 s_l = *reinterpret_cast(&slw[w]); + const __nv_bfloat162 c_r = *reinterpret_cast(&crw[w]); + const __nv_bfloat162 s_r = *reinterpret_cast(&srw[w]); + const __nv_bfloat162 x1 = __hfma2(xl, c_l, __hmul2(xr, s_r)); + const __nv_bfloat162 x2 = __hfma2(xr, c_r, __hneg2(__hmul2(xl, s_l))); + const uint32_t a = *reinterpret_cast(&x1); + const uint32_t b = *reinterpret_cast(&x2); + o[2 * w] = __byte_perm(a, b, 0x5410); // x1[lo], x2[lo] + o[2 * w + 1] = __byte_perm(a, b, 0x7632); // x1[hi], x2[hi] + } +} + +// --------------------------------------------------------------------------------------------- +// Forward: query / key-value RoPE + MXFP8, attention scale layout. +// --------------------------------------------------------------------------------------------- + +__device__ __forceinline__ void cp_async16(const uint32_t smem_addr, const void *gptr) { + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(smem_addr), "l"(gptr) + : "memory"); +} +__device__ __forceinline__ void cp_async_commit() { + asm volatile("cp.async.commit_group;\n" ::: "memory"); +} +template +__device__ __forceinline__ void cp_async_wait() { + asm volatile("cp.async.wait_group %0;\n" ::"n"(N) : "memory"); +} + +// Stores the data of a quantized unit of head column `column` (= batch * heads + head) with head +// dimension D, starting at head dimension d0, tokens t0 .. t0 + 31. +__device__ __forceinline__ void store_data(const QuantizedUnit &o, uint8_t *rowwise_data, + uint8_t *columnwise_data, const uint32_t D, + const uint32_t t0, const uint32_t column, + const uint32_t num_columns, const uint32_t d0, + const int lane) { + const uint32_t r = lane >> 3, k = lane & 7; + const uint32_t row_bytes = num_columns * D; + const uint32_t offset = (t0 + r) * row_bytes + column * D + d0 + 8 * k; + uint8_t *rd = rowwise_data + offset; + uint8_t *cd = columnwise_data + offset; + const uint32_t step = 4 * row_bytes; +#pragma unroll + for (int i = 0; i < 8; ++i) { + *reinterpret_cast(rd) = o.row[i]; + *reinterpret_cast(cd) = o.col[i]; + rd += step; + cd += step; + } +} + +// Places a unit's scales (token block j of the tile, head dimension d0) in the tile's swizzled +// scale blocks of one head, in shared memory: +// row-wise: [D / 128 blocks][128 words]; word (s % 32) * 4 + s / 32 holds the scales of the +// 4 consecutive 32-element blocks of token s; +// column-wise: [ceil(D / 128) blocks][128 words]; word (n % 32) * 4 + (n % 128) / 32 of block +// n / 128 holds the scales of the 4 token blocks of column n. +__device__ __forceinline__ void stage_scales(const QuantizedUnit &o, uint8_t *rowwise_tile, + uint8_t *columnwise_tile, const uint32_t j, + const uint32_t d0, const int lane) { + const uint32_t r = lane >> 3, k = lane & 7; + const uint32_t block = d0 / 32; // Even. + if (k == 0) { + uint8_t *base = rowwise_tile + (block / 4) * 512 + (block % 4) + 4 * j; +#pragma unroll + for (int i = 0; i < 8; ++i) { + *reinterpret_cast(base + 16 * (4 * i + r)) = static_cast(o.rs[i]); + } + } + if (r == 0) { + const uint32_t bytes[2] = {o.cs.x, o.cs.y}; +#pragma unroll + for (int jj = 0; jj < 8; ++jj) { + const uint32_t n = d0 + 8 * k + jj; + columnwise_tile[(n / 128) * 512 + ((n % 32) * 4 + (n % 128) / 32) * 4 + j] = + static_cast(bytes[jj / 4] >> (8 * (jj % 4))); + } + } +} + +// One MXFP8 output in the attention layout: data [s, b * h * D] (SBHD), scaling factors permuted +// to BHSD and swizzled: row-wise [b * h * s, ceil(D / 32 / 4) * 4], column-wise +// [b * h * s / 32, ceil(D / 128) * 128]. +struct AttnOutput { + uint8_t *rowwise_data, *columnwise_data, *rowwise_scale_inv, *columnwise_scale_inv; +}; + +struct RopeParams { + const uint8_t *src; // Query [s, b, h, 192] or key/value [s, b, h, 256], BF16. + const uint8_t *k_rope; // Key/value: s * b rows of 64 BF16, row (t * b + batch). + uint32_t k_rope_row_bytes; // Bytes between consecutive k_rope rows. + const uint16_t *cos, *sin; // [>= s, 64] BF16. + AttnOutput out0, out1; // Query; or key, value. + uint32_t num_columns, heads, seq_len; // b * h, h, s. +}; + +constexpr int kStages = 2; // Load ring stages per warp. +constexpr int kWarps = 8; +constexpr int kThreads = kWarps * THREADS_PER_WARP; +constexpr int kHeadsPerTile = 2; // Heads of one batch entry per tile. +constexpr int kStageBytes = 32 * 128; // One unit of BF16 input: 32 tokens x 64 columns. + +// Shared memory of a CTA: the warps' load rings, then two sets of scale tiles (query: 1 + 1 KB per +// head; key/value: key 1 + 1 KB, value 0.5 + 0.5 KB per head). +template +struct RopeLayout { + static constexpr int kRing = kWarps * kStages * kStageBytes; + static constexpr int kScales = kKV ? kHeadsPerTile * 3072 : kHeadsPerTile * 2048; + static constexpr int kBytes = kRing + 2 * kScales; + // Units per tile (4 token blocks): query 3 per head; key/value 4 per head + 1 RoPE unit + // (the shared key RoPE dimensions) per token block. + static constexpr int kUnits = kKV ? 4 * kHeadsPerTile * 4 + 4 : 4 * kHeadsPerTile * 3; + static constexpr int kIters = (kUnits + kWarps - 1) / kWarps; +}; + +// Persistent: CTA x handles tiles x, x + gridDim.x, ... (tile = token tile * groups + head +// group); each warp streams its units of consecutive tiles through its cp.async ring, so the +// loads of the next tile overlap the end of the current one. +template +__global__ void __launch_bounds__(kThreads) fused_mla_rope_mxfp8_attn_kernel(const RopeParams p) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + using L = RopeLayout; + extern __shared__ __align__(16) uint8_t smem[]; + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const uint32_t r = lane >> 3, k = lane & 7; + uint8_t *ring = smem + warp * kStages * kStageBytes; + uint8_t *scales_base = smem + L::kRing; + const uint32_t ring_addr = + static_cast(__cvta_generic_to_shared(ring)) + r * 128 + 16 * k; + const uint32_t groups = p.num_columns / kHeadsPerTile; + const uint32_t token_tiles = p.seq_len / kTileTokens; + const uint32_t num_tiles = groups * token_tiles; + const uint32_t my_tiles = + blockIdx.x < num_tiles ? (num_tiles - 1 - blockIdx.x) / gridDim.x + 1 : 0; + const uint32_t num_steps = my_tiles * L::kIters; + + for (int i = threadIdx.x; i < 2 * L::kScales / 16; i += kThreads) { + reinterpret_cast(scales_base)[i] = make_uint4(0, 0, 0, 0); + } + + // Unit u of a tile -> (token block j, head hh, part). Query parts: 64-column pieces of the + // head; key/value parts 0, 1: key without RoPE, 2, 3: value, 4: the key RoPE dimensions. + auto decode = [&](const uint32_t u, uint32_t &j, uint32_t &hh, uint32_t &part) { + if (kKV) { + if (u >= 16 * kHeadsPerTile) { + j = u - 16 * kHeadsPerTile; + hh = 0; + part = 4; + } else { + j = u / (4 * kHeadsPerTile); + const uint32_t rem = u - j * 4 * kHeadsPerTile; + hh = rem >> 2; + part = rem & 3; + } + } else { + j = u / (3 * kHeadsPerTile); + const uint32_t rem = u - j * 3 * kHeadsPerTile; + hh = rem / 3; + part = rem - hh * 3; + } + }; + // Issues the loads of this warp's step n into ring stage `stage`. + auto issue = [&](const uint32_t n, const int stage) { + const uint32_t ti = n / L::kIters; + const uint32_t tile = blockIdx.x + ti * gridDim.x; + const uint32_t u = warp + (n - ti * L::kIters) * kWarps; + if (n < num_steps && u < L::kUnits) { + uint32_t j, hh, part; + decode(u, j, hh, part); + const uint32_t tt = tile / groups, col0 = (tile - tt * groups) * kHeadsPerTile; + const uint32_t t = (tt * 4 + j) * 32 + r, column = col0 + hh; + const uint8_t *src; + uint32_t stride; + if (kKV && part == 4) { + const uint32_t batches = p.num_columns / p.heads, batch = col0 / p.heads; + stride = batches * p.k_rope_row_bytes; + src = p.k_rope + static_cast(t * batches + batch) * p.k_rope_row_bytes; + } else { + constexpr uint32_t kRowBytes = (kKV ? kKVDim : kQKDim) * 2; + stride = p.num_columns * kRowBytes; + src = p.src + static_cast(t * p.num_columns + column) * kRowBytes + part * 128; + } + src += 16 * k; +#pragma unroll + for (int i = 0; i < 8; ++i) { + cp_async16(ring_addr + stage * kStageBytes + i * 512, src); + src += 4 * stride; + } + } + cp_async_commit(); + }; + +#pragma unroll + for (int s = 0; s < kStages - 1; ++s) issue(s, s); + __syncthreads(); // Scale tiles zeroed. + int stage = 0; + const uint32_t mc4 = p.num_columns * p.seq_len / kTileTokens; +#pragma unroll 1 + for (uint32_t ti = 0; ti < my_tiles; ++ti) { + const uint32_t tile = blockIdx.x + ti * gridDim.x; + const uint32_t tt = tile / groups, col0 = (tile - tt * groups) * kHeadsPerTile; + uint8_t *scales = scales_base + (ti & 1) * L::kScales; +#pragma unroll 1 + for (int it = 0; it < L::kIters; ++it) { + const uint32_t n = ti * L::kIters + it; + const uint32_t u = warp + it * kWarps; + issue(n + kStages - 1, (stage + kStages - 1) % kStages); + cp_async_wait(); + __syncwarp(); + const bool valid = u < L::kUnits; + uint32_t j, hh, part; + decode(valid ? u : 0, j, hh, part); + const bool rope = kKV ? part == 4 : part == 2; + const uint32_t t0 = (tt * 4 + j) * 32; + const uint8_t *staged = ring + stage * kStageBytes + r * 128; + uint32_t v[8][4]; + if (!rope) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint4 x = *reinterpret_cast(staged + i * 512 + 16 * k); + v[i][0] = x.x; + v[i][1] = x.y; + v[i][2] = x.z; + v[i][3] = x.w; + } + } else { + // Output chunk k: left (k < 4) or right half; rotation pairs 8 * (k % 4) .. + 7, which + // are input chunks 2 * (k % 4) and 2 * (k % 4) + 1. + const uint32_t g = k & 3; + const bool right = k >= 4; + const uint16_t *c = p.cos + (t0 + r) * kRopeDim + (right ? 32 : 0) + 8 * g; + const uint16_t *sn = p.sin + (t0 + r) * kRopeDim + (right ? 32 : 0) + 8 * g; +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint8_t *row = staged + i * 512 + 32 * g; + const uint4 x_lo = *reinterpret_cast(row); + const uint4 x_hi = *reinterpret_cast(row + 16); + const uint4 ca = __ldg(reinterpret_cast(c + i * 4 * kRopeDim)); + const uint4 sa = __ldg(reinterpret_cast(sn + i * 4 * kRopeDim)); + rope8(x_lo, x_hi, ca, sa, right, v[i]); + } + } + __syncwarp(); + QuantizedUnit o; + quantize(v, o); + if (valid) { + if (kKV) { + if (part == 4) { + for (uint32_t h = 0; h < kHeadsPerTile; ++h) { + store_data(o, p.out0.rowwise_data, p.out0.columnwise_data, kQKDim, t0, col0 + h, + p.num_columns, kNopeDim, lane); + stage_scales(o, scales + h * 1024, scales + kHeadsPerTile * 1024 + h * 1024, j, + kNopeDim, lane); + } + } else if (part < 2) { + store_data(o, p.out0.rowwise_data, p.out0.columnwise_data, kQKDim, t0, col0 + hh, + p.num_columns, 64 * part, lane); + stage_scales(o, scales + hh * 1024, scales + kHeadsPerTile * 1024 + hh * 1024, j, + 64 * part, lane); + } else { + store_data(o, p.out1.rowwise_data, p.out1.columnwise_data, kValueDim, t0, col0 + hh, + p.num_columns, 64 * (part - 2), lane); + stage_scales(o, scales + kHeadsPerTile * 2048 + hh * 512, + scales + kHeadsPerTile * 2560 + hh * 512, j, 64 * (part - 2), lane); + } + } else { + store_data(o, p.out0.rowwise_data, p.out0.columnwise_data, kQKDim, t0, col0 + hh, + p.num_columns, 64 * part, lane); + stage_scales(o, scales + hh * 1024, scales + kHeadsPerTile * 1024 + hh * 1024, j, + 64 * part, lane); + } + } + stage = (stage + 1) % kStages; + } + __syncthreads(); // The tile's scales are staged. + // Scale tiles -> global, then zeroed for tile ti + 2. Row-wise block b of column c: + // ((c * s / 128 + tt) * (padded blocks / 4) + b) * 512 bytes. Column-wise block b of column + // c: (b * c_total * s / 128 + c * s / 128 + tt) * 512 bytes. + constexpr int kBlocksPerHead = kKV ? 6 : 4; // Query/key: 2 row-wise + 2 column-wise; + // value: 1 + 1. + for (int task = threadIdx.x; task < kHeadsPerTile * kBlocksPerHead * 32; task += kThreads) { + const int vec = task & 31, blk = task >> 5; + const int hh = blk / kBlocksPerHead, b = blk % kBlocksPerHead; + const uint32_t c = col0 + hh; + uint8_t *src_block; + uint8_t *dst; + if (b < 2) { + src_block = scales + hh * 1024 + b * 512; + dst = p.out0.rowwise_scale_inv + static_cast((c * token_tiles + tt) * 2 + b) * 512; + } else if (b < 4) { + src_block = scales + kHeadsPerTile * 1024 + hh * 1024 + (b - 2) * 512; + dst = p.out0.columnwise_scale_inv + + static_cast((b - 2) * mc4 + c * token_tiles + tt) * 512; + } else if (b == 4) { + src_block = scales + kHeadsPerTile * 2048 + hh * 512; + dst = p.out1.rowwise_scale_inv + static_cast(c * token_tiles + tt) * 512; + } else { + src_block = scales + kHeadsPerTile * 2560 + hh * 512; + dst = p.out1.columnwise_scale_inv + static_cast(c * token_tiles + tt) * 512; + } + uint4 *sv = reinterpret_cast(src_block) + vec; + reinterpret_cast(dst)[vec] = *sv; + *sv = make_uint4(0, 0, 0, 0); + } + } + cp_async_wait<0>(); +#else + NVTE_DEVICE_THREAD0_ERROR("Fused MLA RoPE with MXFP8 output requires SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +// --------------------------------------------------------------------------------------------- +// Backward. +// --------------------------------------------------------------------------------------------- + +// Query RoPE backward in place on dq [rows, 192] (rows = s * b * h). Lane m of a row's 8 lanes +// reads rotation pairs 4m .. 4m + 3 (8 bytes of each half) and writes one 16-byte chunk; a warp +// covers 4 rows at a time, kRowsPerLane times. +constexpr int kQBwdRowsPerLane = 4; +constexpr int kQBwdThreads = 256; + +__global__ void __launch_bounds__(kQBwdThreads) + fused_mla_q_rope_bwd_kernel(uint16_t *__restrict__ dq, const uint16_t *__restrict__ cos, + const uint16_t *__restrict__ sin, const uint32_t rows, + const uint32_t rows_per_token) { + constexpr int NR = kQBwdRowsPerLane; + const uint32_t lane = threadIdx.x & 31, r = lane >> 3, m = lane & 7; + const uint32_t warp = (blockIdx.x * blockDim.x + threadIdx.x) >> 5; + const uint32_t row0 = warp * (4 * NR) + r; + uint2 left[NR], right[NR], cl[NR], sl[NR], cr[NR], sr[NR]; +#pragma unroll + for (int i = 0; i < NR; ++i) { + const uint32_t row = row0 + 4 * i; + if (row < rows) { + const uint16_t *base = dq + static_cast(row) * kQKDim + kNopeDim; + left[i] = *reinterpret_cast(base + 4 * m); + right[i] = *reinterpret_cast(base + 32 + 4 * m); + const uint32_t t = row / rows_per_token; + cl[i] = __ldg(reinterpret_cast(cos + t * kRopeDim + 4 * m)); + sl[i] = __ldg(reinterpret_cast(sin + t * kRopeDim + 4 * m)); + cr[i] = __ldg(reinterpret_cast(cos + t * kRopeDim + 32 + 4 * m)); + sr[i] = __ldg(reinterpret_cast(sin + t * kRopeDim + 32 + 4 * m)); + } + } +#pragma unroll + for (int i = 0; i < NR; ++i) { + const uint32_t row = row0 + 4 * i; + if (row < rows) { + uint32_t o[4]; + rope_bwd4(left[i], right[i], cl[i], sl[i], cr[i], sr[i], o); + *reinterpret_cast(dq + static_cast(row) * kQKDim + kNopeDim + 8 * m) = + make_uint4(o[0], o[1], o[2], o[3]); + } + } +} + +__device__ __forceinline__ uint4 ld_stream(const void *ptr) { + uint4 v; + asm volatile("ld.global.nc.L1::no_allocate.v4.u32 {%0,%1,%2,%3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(ptr)); + return v; +} + +__device__ __forceinline__ uint2 ld_stream2(const void *ptr) { + uint2 v; + asm volatile("ld.global.nc.L1::no_allocate.v2.u32 {%0,%1}, [%2];" + : "=r"(v.x), "=r"(v.y) + : "l"(ptr)); + return v; +} + +// One MXFP8 output [T, N] with GEMM-swizzled scaling factors (row-wise [T, N / 32], column-wise +// [T / 32, N]), T and N multiples of 128. +struct GradOutput { + uint8_t *rowwise_data, *columnwise_data, *rowwise_scale_inv, *columnwise_scale_inv; +}; + +struct GradParams { + const uint8_t *a; // Query: dq [T, h, 192]; key/value: dk [T, h, 192]. + const uint8_t *b; // Key/value: dv [T, h, 128]. + const uint16_t *cos, *sin; // Query: [>= s, 64] BF16. + GradOutput out; + uint32_t tokens, columns, heads, batch; // T = s * b, N, h, b. +}; + +// CTA = 128 tokens x 128 columns, 8 warps; warp w = token block w / 2 (32 tokens) x column half +// w % 2 (64 columns). A CTA's row-wise and column-wise scaling factors are each exactly one +// 512-byte swizzled block. kKV = false: query (RoPE backward on head dimensions 128..191); +// kKV = true: key/value split backward (d_kv = [dk[:128] | dv] per head). +template +__global__ void __launch_bounds__(kThreads) fused_mla_rope_bwd_mxfp8_kernel(const GradParams p) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + __shared__ __align__(16) uint8_t sc[1024]; // [0, 512): row-wise block, [512, 1024): column-wise. + const int warp = threadIdx.x >> 5, lane = threadIdx.x & 31; + const uint32_t r = lane >> 3, k = lane & 7; + const uint32_t j = warp >> 1, half = warp & 1; + const uint32_t t0 = blockIdx.y * kTileTokens + j * 32; // First token of the unit. + const uint32_t c0 = blockIdx.x * 128 + half * 64; // First column of the unit. + + uint32_t v[8][4]; + if (!kKV) { + const uint32_t d0 = c0 % kQKDim; + const uint8_t *src = p.a + (static_cast(t0 + r) * p.columns + c0) * 2; + const size_t step = static_cast(4) * p.columns * 2; + if (d0 != kNopeDim) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint4 x = ld_stream(src + i * step + 16 * k); + v[i][0] = x.x; + v[i][1] = x.y; + v[i][2] = x.z; + v[i][3] = x.w; + } + } else { + uint2 left[8], right[8]; +#pragma unroll + for (int i = 0; i < 8; ++i) { + left[i] = ld_stream2(src + i * step + 8 * k); + right[i] = ld_stream2(src + i * step + 64 + 8 * k); + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint32_t pos = (t0 + 4 * i + r) / p.batch; + const uint2 cl = __ldg(reinterpret_cast(p.cos + pos * kRopeDim + 4 * k)); + const uint2 sl = __ldg(reinterpret_cast(p.sin + pos * kRopeDim + 4 * k)); + const uint2 cr = + __ldg(reinterpret_cast(p.cos + pos * kRopeDim + 32 + 4 * k)); + const uint2 sr = + __ldg(reinterpret_cast(p.sin + pos * kRopeDim + 32 + 4 * k)); + rope_bwd4(left[i], right[i], cl, sl, cr, sr, v[i]); + } + } + } else { + const uint32_t h = c0 / kKVDim, d0 = c0 % kKVDim; + const uint8_t *src; + size_t step; + if (d0 < kNopeDim) { + src = p.a + ((static_cast(t0 + r) * p.heads + h) * kQKDim + d0) * 2; + step = static_cast(4) * p.heads * kQKDim * 2; + } else { + src = p.b + ((static_cast(t0 + r) * p.heads + h) * kValueDim + d0 - kNopeDim) * 2; + step = static_cast(4) * p.heads * kValueDim * 2; + } +#pragma unroll + for (int i = 0; i < 8; ++i) { + const uint4 x = ld_stream(src + i * step + 16 * k); + v[i][0] = x.x; + v[i][1] = x.y; + v[i][2] = x.z; + v[i][3] = x.w; + } + } + + QuantizedUnit o; + quantize(v, o); + + { + const size_t offset = static_cast(t0 + r) * p.columns + c0 + 8 * k; + const size_t step = static_cast(4) * p.columns; +#pragma unroll + for (int i = 0; i < 8; ++i) { + *reinterpret_cast(p.out.rowwise_data + offset + i * step) = o.row[i]; + *reinterpret_cast(p.out.columnwise_data + offset + i * step) = o.col[i]; + } + } + // Scales -> shared memory, swizzled 512-byte blocks: + // row-wise: token tl < 128, scale column c < 4: byte (tl % 32) * 16 + (tl / 32) * 4 + c; + // column-wise: column n < 128, token block jb < 4: byte (n % 32) * 16 + (n / 32) * 4 + jb. + if (k == 0) { +#pragma unroll + for (int i = 0; i < 8; ++i) { + *reinterpret_cast(sc + (4 * i + r) * 16 + j * 4 + half * 2) = + static_cast(o.rs[i]); + } + } + if (r == 0) { + const uint32_t bytes[2] = {o.cs.x, o.cs.y}; +#pragma unroll + for (int jj = 0; jj < 8; ++jj) { + const uint32_t n = half * 64 + 8 * k + jj; + sc[512 + (n % 32) * 16 + (n / 32) * 4 + j] = + static_cast(bytes[jj / 4] >> (8 * (jj % 4))); + } + } + __syncthreads(); + if (threadIdx.x < 64) { + const int vec = threadIdx.x & 31; + const uint4 x = reinterpret_cast(sc + (threadIdx.x >> 5) * 512)[vec]; + uint8_t *dst = + (threadIdx.x < 32) + ? p.out.rowwise_scale_inv + + (static_cast(blockIdx.y) * (p.columns / 128) + blockIdx.x) * 512 + : p.out.columnwise_scale_inv + + (static_cast(blockIdx.x) * (p.tokens / kTileTokens) + blockIdx.y) * 512; + reinterpret_cast(dst)[vec] = x; + } +#else + NVTE_DEVICE_THREAD0_ERROR("Fused MLA RoPE backward with MXFP8 output requires SM 10.0+."); +#endif // (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +} + +// --------------------------------------------------------------------------------------------- +// Host side. +// --------------------------------------------------------------------------------------------- + +bool is_aligned_to(const void *ptr, const size_t alignment) { + return reinterpret_cast(ptr) % alignment == 0; +} + +void check_bf16_input(const Tensor &t, const char *name) { + NVTE_CHECK(t.data.dtype == DType::kBFloat16, name, " must be BF16."); + NVTE_CHECK(is_aligned_to(t.data.dptr, 16), name, " must be 16-byte aligned."); +} + +void check_cos_sin(const Tensor &cos, const Tensor &sin, const size_t s) { + check_bf16_input(cos, "cos"); + check_bf16_input(sin, "sin"); + NVTE_CHECK(cos.flat_last_dim() == kRopeDim && sin.flat_last_dim() == kRopeDim && + cos.flat_first_dim() >= s && sin.flat_first_dim() >= s, + "cos and sin must have at least ", s, " rows of ", kRopeDim, " elements, got shapes ", + cos.data.shape, " and ", sin.data.shape, "."); +} + +// An MXFP8 output with row-wise and column-wise E4M3 data and GEMM-swizzled scaling factors, +// whose flattened 2D shape is [rows, cols]. +GradOutput check_swizzled_mxfp8_output(const Tensor &t, const char *name, const size_t rows, + const size_t cols) { + NVTE_CHECK(is_mxfp8_scaling(t.scaling_mode), name, " must be an MXFP8 tensor, got ", + to_string(t.scaling_mode), "."); + NVTE_CHECK(t.has_data() && t.has_columnwise_data(), name, + " needs both row-wise and column-wise data."); + NVTE_CHECK(t.data.dtype == DType::kFloat8E4M3 && t.columnwise_data.dtype == DType::kFloat8E4M3, + name, " must have E4M3 data."); + NVTE_CHECK(t.with_gemm_swizzled_scales, name, " must have GEMM-swizzled scaling factors."); + const auto [first_dim, last_dim] = t.flat_2d_dims(); + NVTE_CHECK(first_dim == rows && last_dim == cols, name, " must have flattened shape [", rows, + ", ", cols, "], got ", t.shape(), "."); + CheckOutputTensor(t, name); + GradOutput out{static_cast(t.data.dptr), + static_cast(t.columnwise_data.dptr), + static_cast(t.scale_inv.dptr), + static_cast(t.columnwise_scale_inv.dptr)}; + NVTE_CHECK(is_aligned_to(out.rowwise_data, 16) && is_aligned_to(out.columnwise_data, 16) && + is_aligned_to(out.rowwise_scale_inv, 16) && + is_aligned_to(out.columnwise_scale_inv, 16), + name, " buffers must be 16-byte aligned."); + return out; +} + +AttnOutput to_attn_output(const GradOutput &o) { + return AttnOutput{o.rowwise_data, o.columnwise_data, o.rowwise_scale_inv, o.columnwise_scale_inv}; +} + +template +void launch_rope_mxfp8_attn(const RopeParams &p, cudaStream_t stream) { + using L = RopeLayout; + const size_t num_tiles = + static_cast(p.num_columns / kHeadsPerTile) * (p.seq_len / kTileTokens); + if (num_tiles == 0) return; + auto kernel = fused_mla_rope_mxfp8_attn_kernel; + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, L::kBytes)); + int ctas_per_sm = 0; + NVTE_CHECK_CUDA( + cudaOccupancyMaxActiveBlocksPerMultiprocessor(&ctas_per_sm, kernel, kThreads, L::kBytes)); + const size_t grid = + std::min(static_cast(cuda::sm_count()) * std::max(ctas_per_sm, 1), num_tiles); + kernel<<>>(p); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Checks of the forward inputs shared by query and key/value; returns (s, b, h). +std::array check_rope_input(const Tensor &x, const char *name, const size_t head_dim) { + check_bf16_input(x, name); + const auto &shape = x.data.shape; + NVTE_CHECK(shape.size() == 4 && shape[3] == head_dim, name, " must have shape [s, b, h, ", + head_dim, "], got ", shape, "."); + const size_t s = shape[0], b = shape[1], h = shape[2]; + NVTE_CHECK(s % kTileTokens == 0, "The sequence length (", s, ") must be a multiple of ", + kTileTokens, "."); + NVTE_CHECK(h % kHeadsPerTile == 0, "The number of heads (", h, ") must be even."); + // 32-bit offsets in the kernel. + NVTE_CHECK(s * b * h * kKVDim * 2 < (static_cast(1) << 31), name, " shape ", shape, + " is too large."); + return {s, b, h}; +} + +template +void launch_rope_bwd_mxfp8(const GradParams &p, cudaStream_t stream) { + if (p.tokens == 0) return; + const dim3 grid(p.columns / 128, p.tokens / kTileTokens); + fused_mla_rope_bwd_mxfp8_kernel<<>>(p); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +// Checks of the backward gradients; returns (s, b, h). +std::array check_grad_input(const Tensor &x, const char *name, const size_t head_dim) { + check_bf16_input(x, name); + const auto &shape = x.data.shape; + NVTE_CHECK(shape.size() == 4 && shape[3] == head_dim, name, " must have shape [s, b, h, ", + head_dim, "], got ", shape, "."); + const size_t s = shape[0], b = shape[1], h = shape[2]; + NVTE_CHECK((s * b) % kTileTokens == 0, "s * b (", s * b, ") must be a multiple of ", kTileTokens, + "."); + NVTE_CHECK(h % 2 == 0, "The number of heads (", h, ") must be even."); + NVTE_CHECK(s * b * h * kKVDim < (static_cast(1) << 31) && (s * b) / kTileTokens <= 65535, + name, " shape ", shape, " is too large."); + return {s, b, h}; +} + +} // namespace + +void fused_mla_q_rope_mxfp8_attn(const Tensor &q, const Tensor &cos, const Tensor &sin, + Tensor *query, cudaStream_t stream) { + NVTE_CHECK(is_supported_by_CC_100(), "Fused MLA query RoPE with MXFP8 output requires SM 10.0+."); + const auto [s, b, h] = check_rope_input(q, "q", kQKDim); + check_cos_sin(cos, sin, s); + RopeParams p{}; + p.src = static_cast(q.data.dptr); + p.cos = static_cast(cos.data.dptr); + p.sin = static_cast(sin.data.dptr); + p.out0 = to_attn_output(check_swizzled_mxfp8_output(*query, "query", s * b * h, kQKDim)); + p.out1 = p.out0; + p.num_columns = static_cast(b * h); + p.heads = static_cast(h); + p.seq_len = static_cast(s); + launch_rope_mxfp8_attn(p, stream); +} + +void fused_mla_kv_rope_mxfp8_attn(const Tensor &kv, const Tensor &k_rope, const int k_rope_stride, + const Tensor &cos, const Tensor &sin, Tensor *key, Tensor *value, + cudaStream_t stream) { + NVTE_CHECK(is_supported_by_CC_100(), + "Fused MLA key/value RoPE with MXFP8 output requires SM 10.0+."); + const auto [s, b, h] = check_rope_input(kv, "kv", kKVDim); + check_bf16_input(k_rope, "k_rope"); + NVTE_CHECK(k_rope.flat_first_dim() == s * b && k_rope.flat_last_dim() == kRopeDim, + "k_rope must have s * b = ", s * b, " rows of ", kRopeDim, " elements, got shape ", + k_rope.data.shape, "."); + NVTE_CHECK(k_rope_stride >= kRopeDim && k_rope_stride % 8 == 0, "The k_rope row stride (", + k_rope_stride, ") must be at least the row size and a multiple of 8 elements."); + NVTE_CHECK(static_cast(k_rope_stride) * 2 * s * b < (static_cast(1) << 31), + "k_rope is too large."); + check_cos_sin(cos, sin, s); + RopeParams p{}; + p.src = static_cast(kv.data.dptr); + p.k_rope = static_cast(k_rope.data.dptr); + p.k_rope_row_bytes = static_cast(k_rope_stride) * 2; + p.cos = static_cast(cos.data.dptr); + p.sin = static_cast(sin.data.dptr); + p.out0 = to_attn_output(check_swizzled_mxfp8_output(*key, "key", s * b * h, kQKDim)); + p.out1 = to_attn_output(check_swizzled_mxfp8_output(*value, "value", s * b * h, kValueDim)); + p.num_columns = static_cast(b * h); + p.heads = static_cast(h); + p.seq_len = static_cast(s); + launch_rope_mxfp8_attn(p, stream); +} + +void fused_mla_q_rope_backward(Tensor *dq, const Tensor &cos, const Tensor &sin, + cudaStream_t stream) { + check_bf16_input(*dq, "dq"); + const auto &shape = dq->data.shape; + NVTE_CHECK(shape.size() == 4 && shape[3] == kQKDim, "dq must have shape [s, b, h, ", kQKDim, + "], got ", shape, "."); + const size_t s = shape[0], b = shape[1], h = shape[2]; + NVTE_CHECK(s * b * h * kQKDim * 2 < (static_cast(1) << 32), "dq shape ", shape, + " is too large."); + check_cos_sin(cos, sin, s); + const size_t rows = s * b * h; + if (rows == 0) return; + constexpr size_t kRowsPerCTA = kQBwdThreads / THREADS_PER_WARP * 4 * kQBwdRowsPerLane; + const size_t grid = DIVUP(rows, kRowsPerCTA); + fused_mla_q_rope_bwd_kernel<<>>( + static_cast(dq->data.dptr), static_cast(cos.data.dptr), + static_cast(sin.data.dptr), static_cast(rows), + static_cast(b * h)); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void fused_mla_q_rope_backward_mxfp8(const Tensor &dq, const Tensor &cos, const Tensor &sin, + Tensor *grad, cudaStream_t stream) { + NVTE_CHECK(is_supported_by_CC_100(), + "Fused MLA query RoPE backward with MXFP8 output requires SM 10.0+."); + const auto [s, b, h] = check_grad_input(dq, "dq", kQKDim); + check_cos_sin(cos, sin, s); + GradParams p{}; + p.a = static_cast(dq.data.dptr); + p.cos = static_cast(cos.data.dptr); + p.sin = static_cast(sin.data.dptr); + p.out = check_swizzled_mxfp8_output(*grad, "grad", s * b, h * kQKDim); + p.tokens = static_cast(s * b); + p.columns = static_cast(h * kQKDim); + p.heads = static_cast(h); + p.batch = static_cast(b); + launch_rope_bwd_mxfp8(p, stream); +} + +void fused_mla_kv_split_backward_mxfp8(const Tensor &dk, const Tensor &dv, Tensor *grad, + cudaStream_t stream) { + NVTE_CHECK(is_supported_by_CC_100(), + "Fused MLA key/value split backward with MXFP8 output requires SM 10.0+."); + const auto [s, b, h] = check_grad_input(dk, "dk", kQKDim); + check_bf16_input(dv, "dv"); + NVTE_CHECK(dv.data.shape == std::vector({s, b, h, static_cast(kValueDim)}), + "dv must have shape [", s, ", ", b, ", ", h, ", ", kValueDim, "], got ", dv.data.shape, + "."); + GradParams p{}; + p.a = static_cast(dk.data.dptr); + p.b = static_cast(dv.data.dptr); + p.out = check_swizzled_mxfp8_output(*grad, "grad", s * b, h * kKVDim); + p.tokens = static_cast(s * b); + p.columns = static_cast(h * kKVDim); + p.heads = static_cast(h); + p.batch = static_cast(b); + launch_rope_bwd_mxfp8(p, stream); +} + +} // namespace transformer_engine + +void nvte_fused_mla_q_rope_mxfp8_attn(const NVTETensor q, const NVTETensor cos, + const NVTETensor sin, NVTETensor query, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_mla_q_rope_mxfp8_attn); + using namespace transformer_engine; + fused_mla_q_rope_mxfp8_attn(*convertNVTETensorCheck(q), *convertNVTETensorCheck(cos), + *convertNVTETensorCheck(sin), convertNVTETensorCheck(query), stream); +} + +void nvte_fused_mla_kv_rope_mxfp8_attn(const NVTETensor kv, const NVTETensor k_rope, + const int k_rope_stride, const NVTETensor cos, + const NVTETensor sin, NVTETensor key, NVTETensor value, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_mla_kv_rope_mxfp8_attn); + using namespace transformer_engine; + fused_mla_kv_rope_mxfp8_attn(*convertNVTETensorCheck(kv), *convertNVTETensorCheck(k_rope), + k_rope_stride, *convertNVTETensorCheck(cos), + *convertNVTETensorCheck(sin), convertNVTETensorCheck(key), + convertNVTETensorCheck(value), stream); +} + +void nvte_fused_mla_q_rope_backward(NVTETensor dq, const NVTETensor cos, const NVTETensor sin, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_mla_q_rope_backward); + using namespace transformer_engine; + fused_mla_q_rope_backward(convertNVTETensorCheck(dq), *convertNVTETensorCheck(cos), + *convertNVTETensorCheck(sin), stream); +} + +void nvte_fused_mla_q_rope_backward_mxfp8(const NVTETensor dq, const NVTETensor cos, + const NVTETensor sin, NVTETensor grad, + cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_mla_q_rope_backward_mxfp8); + using namespace transformer_engine; + fused_mla_q_rope_backward_mxfp8(*convertNVTETensorCheck(dq), *convertNVTETensorCheck(cos), + *convertNVTETensorCheck(sin), convertNVTETensorCheck(grad), + stream); +} + +void nvte_fused_mla_kv_split_backward_mxfp8(const NVTETensor dk, const NVTETensor dv, + NVTETensor grad, cudaStream_t stream) { + NVTE_API_CALL(nvte_fused_mla_kv_split_backward_mxfp8); + using namespace transformer_engine; + fused_mla_kv_split_backward_mxfp8(*convertNVTETensorCheck(dk), *convertNVTETensorCheck(dv), + convertNVTETensorCheck(grad), stream); +} diff --git a/transformer_engine/common/include/transformer_engine/fused_rope.h b/transformer_engine/common/include/transformer_engine/fused_rope.h index f0e4540b7cc..042bb101118 100644 --- a/transformer_engine/common/include/transformer_engine/fused_rope.h +++ b/transformer_engine/common/include/transformer_engine/fused_rope.h @@ -180,6 +180,111 @@ void nvte_fused_mla_kv_rope_mxfp8(const NVTETensor kv, const NVTETensor k_rope, const NVTETensor sin, NVTETensor key, NVTETensor value, cudaStream_t stream); +/*! \brief Apply RoPE to multi-latent attention's query and quantize it to MXFP8 for attention. + * + * Each head of q holds 128 dimensions without RoPE followed by 64 RoPE dimensions with the + * rotation pairs interleaved. The RoPE dimensions are rotated as in + * nvte_fused_mla_kv_rope_mxfp8 (Megatron-LM's MLA RoPE with BF16 cos and sin), and the query is + * quantized row-wise (32 elements of a head) and column-wise (32 tokens) with E4M3 data. + * + * The data are written in SBHD. The scaling factors are written permuted to BHSD and swizzled + * for the GEMM: the layout that MXFP8 fused attention consumes, the bytes that the MXFP8 + * quantizer of the RoPE'd query followed by mxfp8_transpose_swizzle (PyTorch) give. + * + * Requires SM 10.0+, an even number of heads and a sequence length that is a multiple of 128. + * + * \param[in] q Query, [s, b, h, 192] in BF16. + * \param[in] cos Cosines, [max_s, 64] in BF16; row t is used for token t. + * \param[in] sin Sines, same shape and type as cos. + * \param[in,out] query MXFP8 query of shape [s, b, h, 192] with row-wise and column-wise + * data and GEMM-swizzled scaling factors, row-wise [b * h * s, 8] and + * column-wise [b * h * s / 32, 256], rows in (b, h, s) order. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_mla_q_rope_mxfp8_attn(const NVTETensor q, const NVTETensor cos, + const NVTETensor sin, NVTETensor query, cudaStream_t stream); + +/*! \brief nvte_fused_mla_kv_rope_mxfp8 with the scaling factors in the layout of MXFP8 fused + * attention. + * + * Key and value data as nvte_fused_mla_kv_rope_mxfp8 (SBHD); scaling factors permuted to BHSD + * and swizzled for the GEMM, as nvte_fused_mla_q_rope_mxfp8_attn. BF16 cos and sin only. + * + * Requires SM 10.0+, an even number of heads and a sequence length that is a multiple of 128. + * + * \param[in] kv Key/value, [s, b, h, 256] in BF16. + * \param[in] k_rope Key RoPE dimensions shared by the heads, s * b rows + * (token-major) of 64 BF16 values. + * \param[in] k_rope_stride Elements from one row of k_rope to the next, a multiple of 8. + * \param[in] cos Cosines, [max_s, 64] in BF16; row t is used for token t. + * \param[in] sin Sines, same shape and type as cos. + * \param[in,out] key MXFP8 key of shape [s, b, h, 192], scaling factors row-wise + * [b * h * s, 8] and column-wise [b * h * s / 32, 256]. + * \param[in,out] value MXFP8 value of shape [s, b, h, 128], scaling factors row-wise + * [b * h * s, 4] and column-wise [b * h * s / 32, 128]. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_mla_kv_rope_mxfp8_attn(const NVTETensor kv, const NVTETensor k_rope, + const int k_rope_stride, const NVTETensor cos, + const NVTETensor sin, NVTETensor key, NVTETensor value, + cudaStream_t stream); + +/*! \brief Backward of the multi-latent attention query RoPE, in place. + * + * For the RoPE dimensions of each head of dq (left = dq[128 : 160], right = dq[160 : 192]), + * writes the gradient of the interleaved input pairs: + * + * dq[128 + 2i] = left_i * cos[t, i] + right_i * sin[t, 32 + i] + * dq[128 + 2i + 1] = right_i * cos[t, 32 + i] - left_i * sin[t, i] + * + * with the roundings of Megatron-LM's MLA RoPE backward (BF16 multiply, negate and fma). + * + * \param[in,out] dq Query gradient, [s, b, h, 192] in BF16, contiguous. + * \param[in] cos Cosines, [max_s, 64] in BF16. + * \param[in] sin Sines, same shape and type as cos. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_mla_q_rope_backward(NVTETensor dq, const NVTETensor cos, const NVTETensor sin, + cudaStream_t stream); + +/*! \brief Backward of the multi-latent attention query RoPE, quantized to MXFP8. + * + * The result of nvte_fused_mla_q_rope_backward, viewed as [s * b, h * 192] (the gradient of the + * query up-projection output), quantized to MXFP8 row-wise and column-wise with E4M3 data and + * GEMM-swizzled scaling factors, as the MXFP8 quantizer with optimize_for_gemm does. The BF16 + * result is not written; dq is not modified. + * + * Requires SM 10.0+, an even number of heads and s * b a multiple of 128. + * + * \param[in] dq Query gradient, [s, b, h, 192] in BF16, contiguous. + * \param[in] cos Cosines, [max_s, 64] in BF16. + * \param[in] sin Sines, same shape and type as cos. + * \param[in,out] grad MXFP8 tensor of shape [s * b, h * 192] with row-wise and column-wise + * data and GEMM-swizzled scaling factors. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_mla_q_rope_backward_mxfp8(const NVTETensor dq, const NVTETensor cos, + const NVTETensor sin, NVTETensor grad, + cudaStream_t stream); + +/*! \brief Key/value split backward of multi-latent attention, quantized to MXFP8. + * + * The gradient of the key/value up-projection output, d_kv[t, h] = [dk[t, h, :128] | dv[t, h]] + * viewed as [s * b, h * 256], quantized to MXFP8 row-wise and column-wise with E4M3 data and + * GEMM-swizzled scaling factors. The gradient of the key RoPE dimensions (dk[..., 128:]) is not + * computed. + * + * Requires SM 10.0+, an even number of heads and s * b a multiple of 128. + * + * \param[in] dk Key gradient, [s, b, h, 192] in BF16, contiguous. + * \param[in] dv Value gradient, [s, b, h, 128] in BF16, contiguous. + * \param[in,out] grad MXFP8 tensor of shape [s * b, h * 256] with row-wise and column-wise + * data and GEMM-swizzled scaling factors. + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_fused_mla_kv_split_backward_mxfp8(const NVTETensor dk, const NVTETensor dv, + NVTETensor grad, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 2ecf8a60090..b2371f0ed3c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -3372,7 +3372,14 @@ def _mxfp8_tensor_in_src_format(fp8_2d, shape, dtype, quantizer, src_format): def fused_mla_kv_rope_mxfp8( - kv, k_rope, cos, sin, v_head_dim, key_quantizer=None, value_quantizer=None + kv, + k_rope, + cos, + sin, + v_head_dim, + key_quantizer=None, + value_quantizer=None, + transpose_swizzle_scales=False, ): """MLA key/value split, key RoPE and MXFP8 quantization in one kernel. @@ -3400,13 +3407,33 @@ def fused_mla_kv_rope_mxfp8( Value dimensions per head in ``kv``. key_quantizer, value_quantizer : MXFP8Quantizer, optional E4M3 quantizers with row-wise and column-wise usage; created if not given. + transpose_swizzle_scales : bool, default = False + Write the scale_invs permuted to BHSD and GEMM-swizzled, the layout MXFP8 attention + consumes: the result of :func:`mxfp8_transpose_swizzle` applied to the outputs, without + running it (see ``nvte_fused_mla_kv_rope_mxfp8_attn``). Needs BF16 ``cos`` and ``sin`` + and a sequence length that is a multiple of 128. Returns ------- key, value : MXFP8Tensor - [s, b, h, k_dim + emb_dim] and [s, b, h, v_head_dim], data and scale_invs in SBHD - layout; NOT yet BHSD-permuted or swizzled (see :func:`mxfp8_transpose_swizzle`). + [s, b, h, k_dim + emb_dim] and [s, b, h, v_head_dim], data in SBHD layout. Scale_invs in + SBHD layout, NOT yet BHSD-permuted or swizzled (see :func:`mxfp8_transpose_swizzle`), + unless ``transpose_swizzle_scales``. """ + if transpose_swizzle_scales: + key_quantizer, value_quantizer = ( + _check_mla_mxfp8_quantizer(quantizer, "fused_mla_kv_rope_mxfp8") + for quantizer in (key_quantizer, value_quantizer) + ) + key_buffers, value_buffers = tex.fused_mla_kv_rope_mxfp8_attn( + kv, k_rope, cos, sin, v_head_dim + ) + key_shape = tuple(key_buffers[0].shape) + value_shape = tuple(value_buffers[0].shape) + return ( + _swizzled_mxfp8_tensor(key_buffers, key_shape, kv.dtype, key_quantizer), + _swizzled_mxfp8_tensor(value_buffers, value_shape, kv.dtype, value_quantizer), + ) s, b, h, kv_dim = kv.shape if s % (4 * MXFP8_BLOCK_SCALING_SIZE) != 0: raise ValueError( @@ -3439,6 +3466,148 @@ def fused_mla_kv_rope_mxfp8( ) +def _check_mla_mxfp8_quantizer(quantizer, name): + """The given MXFP8 quantizer, or a new E4M3 one, checked for the fused MLA kernels.""" + if quantizer is None: + return MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=True) + if not isinstance(quantizer, MXFP8Quantizer): + raise TypeError(f"{name} requires MXFP8 quantizers, got {type(quantizer).__name__}.") + if quantizer.dtype != tex.DType.kFloat8E4M3: + raise ValueError(f"{name} only produces E4M3 data, got {quantizer.dtype}.") + if getattr(quantizer, "with_2d_quantization", False): + raise ValueError(f"{name} does not support 2D MXFP8 quantization.") + return quantizer + + +def _swizzled_mxfp8_tensor(buffers, shape, dtype, quantizer): + """MXFP8Tensor of row-wise data, row-wise scale_inv, column-wise data and column-wise + scale_inv buffers whose scale_invs are GEMM-swizzled.""" + rowwise_data, rowwise_scale_inv, columnwise_data, columnwise_scale_inv = buffers + return MXFP8Tensor( + shape=shape, + dtype=dtype, + rowwise_data=rowwise_data.view(shape), + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data.view(shape), + columnwise_scale_inv=columnwise_scale_inv, + quantizer=quantizer, + requires_grad=False, + fp8_dtype=tex.DType.kFloat8E4M3, + with_gemm_swizzled_scales=True, + ) + + +def fused_mla_q_rope_mxfp8(q, cos, sin, quantizer=None): + """MLA query RoPE and MXFP8 quantization for attention in one kernel. + + Applies Megatron-LM's MLA RoPE (interleaved rotation pairs in, the two halves out, BF16 + roundings) to the last 64 dimensions of each head and quantizes the query row-wise and + column-wise, as :func:`mxfp8_quantize_only` (``src_format="sbhd"``) followed by + :func:`mxfp8_transpose_swizzle` does for the RoPE'd query, with byte-identical results: + data in SBHD, scale_invs permuted to BHSD and GEMM-swizzled. See + ``nvte_fused_mla_q_rope_mxfp8_attn``. Supports DeepSeek-V3's query head dimension (128 + without RoPE + 64 RoPE) on SM 10.0+. + + Parameters + ---------- + q : torch.Tensor + [s, b, h, 192] BF16. ``s`` must be a multiple of 128 and ``h`` even. + cos, sin : torch.Tensor + [max_s, 1, 1, 64] or [max_s, 64] BF16; row t is used for token t. + quantizer : MXFP8Quantizer, optional + E4M3 quantizer with row-wise and column-wise usage; created if not given. + + Returns + ------- + query : MXFP8Tensor + [s, b, h, 192], ready for MXFP8 attention. + """ + quantizer = _check_mla_mxfp8_quantizer(quantizer, "fused_mla_q_rope_mxfp8") + buffers = tex.fused_mla_q_rope_mxfp8_attn(q, cos, sin) + return _swizzled_mxfp8_tensor(buffers, tuple(q.shape), q.dtype, quantizer) + + +def fused_mla_q_rope_backward_(dq, cos, sin): + """Backward of the MLA query RoPE of :func:`fused_mla_q_rope_mxfp8`, in place. + + The same BF16 operations as Megatron-LM's MLA query RoPE backward, so the same bytes. See + ``nvte_fused_mla_q_rope_backward``. + + Parameters + ---------- + dq : torch.Tensor + [s, b, h, 192] BF16, contiguous; its last 64 dimensions per head are overwritten. + cos, sin : torch.Tensor + As for the forward. + + Returns + ------- + dq : torch.Tensor + """ + return tex.fused_mla_q_rope_backward_(dq, cos, sin) + + +def fused_mla_q_rope_backward_mxfp8(dq, cos, sin, quantizer=None): + """MLA query RoPE backward fused with the MXFP8 quantization of its result. + + Returns what an MXFP8 quantizer with ``optimize_for_gemm`` gives for + ``fused_mla_q_rope_backward_(dq, cos, sin).view(s, b, h * 192)``, the gradient of the query + up-projection output: row-wise and column-wise E4M3 data and GEMM-swizzled scale_invs, + byte-identical, without writing the BF16 gradient (``dq`` is not modified). See + ``nvte_fused_mla_q_rope_backward_mxfp8``. + + Parameters + ---------- + dq : torch.Tensor + [s, b, h, 192] BF16; ``s * b`` must be a multiple of 128 and ``h`` even. + cos, sin : torch.Tensor + As for the forward. + quantizer : MXFP8Quantizer, optional + E4M3 quantizer to attach to the result; created if not given. + + Returns + ------- + grad : MXFP8Tensor + [s, b, h * 192]. + """ + quantizer = _check_mla_mxfp8_quantizer(quantizer, "fused_mla_q_rope_backward_mxfp8") + s, b, h, d = dq.shape + buffers = tex.fused_mla_q_rope_backward_mxfp8(dq, cos, sin) + return _swizzled_mxfp8_tensor(buffers, (s, b, h * d), dq.dtype, quantizer) + + +def fused_mla_kv_split_backward_mxfp8(dk, dv, k_dim, quantizer=None): + """MLA key/value split backward fused with the MXFP8 quantization of its result. + + The gradient of the key/value up-projection output, ``[dk[..., :k_dim] | dv]`` per head + viewed as [s, b, h * (k_dim + v_dim)], quantized as an MXFP8 quantizer with + ``optimize_for_gemm`` does (row-wise and column-wise E4M3 data, GEMM-swizzled scale_invs), + byte-identical, without writing it in BF16. The gradient of the key RoPE dimensions + (``dk[..., k_dim:]`` summed over heads) is not computed. See + ``nvte_fused_mla_kv_split_backward_mxfp8``. + + Parameters + ---------- + dk : torch.Tensor + [s, b, h, 192] BF16; ``s * b`` must be a multiple of 128 and ``h`` even. + dv : torch.Tensor + [s, b, h, 128] BF16. + k_dim : int + Key dimensions without RoPE (128). + quantizer : MXFP8Quantizer, optional + E4M3 quantizer to attach to the result; created if not given. + + Returns + ------- + grad : MXFP8Tensor + [s, b, h * (k_dim + v_dim)]. + """ + quantizer = _check_mla_mxfp8_quantizer(quantizer, "fused_mla_kv_split_backward_mxfp8") + s, b, h, _ = dk.shape + buffers = tex.fused_mla_kv_split_backward_mxfp8(dk, dv, k_dim) + return _swizzled_mxfp8_tensor(buffers, (s, b, h * (k_dim + dv.shape[-1])), dk.dtype, quantizer) + + def mxfp8_transpose_swizzle(fp8_tensors, src_format): """Phase 2 of mxfp8_quantize_fast_path: batched BHSD-transpose + GEMM-swizzle. diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 7bcd8430ce1..130576c8d6e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -583,6 +583,22 @@ std::tuple fused_mla_kv_rope_mxfp8( const at::Tensor &kv, const at::Tensor &k_rope, const at::Tensor &cos, const at::Tensor &sin, const int64_t v_head_dim, py::handle key_quantizer, py::handle value_quantizer); +std::vector fused_mla_q_rope_mxfp8_attn(const at::Tensor &q, const at::Tensor &cos, + const at::Tensor &sin); + +std::tuple, std::vector> fused_mla_kv_rope_mxfp8_attn( + const at::Tensor &kv, const at::Tensor &k_rope, const at::Tensor &cos, const at::Tensor &sin, + const int64_t v_head_dim); + +at::Tensor fused_mla_q_rope_backward_(at::Tensor dq, const at::Tensor &cos, const at::Tensor &sin); + +std::vector fused_mla_q_rope_backward_mxfp8(const at::Tensor &dq, const at::Tensor &cos, + const at::Tensor &sin); + +std::vector fused_mla_kv_split_backward_mxfp8(const at::Tensor &dk, + const at::Tensor &dv, + const int64_t k_dim); + /*************************************************************************************************** * Miscellaneous **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp index aea334972c2..5656f7cde19 100644 --- a/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp +++ b/transformer_engine/pytorch/csrc/extensions/apply_rope.cpp @@ -363,4 +363,178 @@ std::tuple fused_mla_kv_rope_mxfp8( return {key_py, value_py}; } +namespace { + +// Buffers of an MXFP8 tensor: row-wise data, row-wise scale_inv, column-wise data, column-wise +// scale_inv (uint8). +std::vector empty_mxfp8_buffers(const std::vector &data_shape, + const std::vector &rowwise_scale_shape, + const std::vector &columnwise_scale_shape, + const at::Device &device) { + const auto opts = at::TensorOptions().dtype(at::kByte).device(device); + return {at::empty(data_shape, opts), at::empty(rowwise_scale_shape, opts), + at::empty(data_shape, opts), at::empty(columnwise_scale_shape, opts)}; +} + +// TE view of MXFP8 buffers (empty_mxfp8_buffers) with GEMM-swizzled scaling factors. +TensorWrapper make_swizzled_mxfp8_tensor(std::vector &buffers) { + auto shape_of = [](const at::Tensor &t) { + std::vector shape; + for (const auto d : t.sizes()) shape.push_back(static_cast(d)); + return shape; + }; + TensorWrapper out(NVTE_MXFP8_1D_SCALING); + out.set_rowwise_data(buffers[0].data_ptr(), DType::kFloat8E4M3, shape_of(buffers[0])); + out.set_rowwise_scale_inv(buffers[1].data_ptr(), DType::kFloat8E8M0, shape_of(buffers[1])); + out.set_columnwise_data(buffers[2].data_ptr(), DType::kFloat8E4M3, shape_of(buffers[2])); + out.set_columnwise_scale_inv(buffers[3].data_ptr(), DType::kFloat8E8M0, shape_of(buffers[3])); + out.set_with_gemm_swizzled_scales(true); + return out; +} + +// MXFP8 buffers of an attention input [s, b, h, d] with BHSD-permuted, GEMM-swizzled scaling +// factors (as mxfp8_transpose_swizzle leaves them). +std::vector empty_mxfp8_attn_buffers(const int64_t s, const int64_t b, const int64_t h, + const int64_t d, const at::Device &device) { + const int64_t rows = b * h * s; + return empty_mxfp8_buffers({s, b, h, d}, {rows, (d / 32 + 3) / 4 * 4}, + {rows / 32, (d + 127) / 128 * 128}, device); +} + +void check_mla_rope_inputs(const at::Tensor &x, const char *name, + const std::vector &others) { + TORCH_CHECK(x.is_cuda(), "expected ", name, " on a CUDA device"); + for (const auto &t : others) { + TORCH_CHECK(t.device() == x.device(), "expected all inputs on the device of ", name); + } + TORCH_CHECK(x.dim() == 4, "expected ", name, " of shape [s, b, h, d], got ", x.sizes()); +} + +} // namespace + +std::vector fused_mla_q_rope_mxfp8_attn(const at::Tensor &q, const at::Tensor &cos, + const at::Tensor &sin) { + check_mla_rope_inputs(q, "q", {cos, sin}); + at::cuda::CUDAGuard device_guard(q.device()); + const auto q_c = q.contiguous(); + const auto cos_c = cos.contiguous(); + const auto sin_c = sin.contiguous(); + auto out = empty_mxfp8_attn_buffers(q.size(0), q.size(1), q.size(2), q.size(3), q.device()); + auto q_cu = makeTransformerEngineTensor(q_c); + auto cos_cu = makeTransformerEngineTensor(cos_c); + auto sin_cu = makeTransformerEngineTensor(sin_c); + auto out_cu = make_swizzled_mxfp8_tensor(out); + NVTE_SCOPED_GIL_RELEASE({ + nvte_fused_mla_q_rope_mxfp8_attn(q_cu.data(), cos_cu.data(), sin_cu.data(), out_cu.data(), + at::cuda::getCurrentCUDAStream()); + }); + return out; +} + +std::tuple, std::vector> fused_mla_kv_rope_mxfp8_attn( + const at::Tensor &kv, const at::Tensor &k_rope, const at::Tensor &cos, const at::Tensor &sin, + const int64_t v_head_dim) { + check_mla_rope_inputs(kv, "kv", {k_rope, cos, sin}); + at::cuda::CUDAGuard device_guard(kv.device()); + const int64_t s = kv.size(0), b = kv.size(1), h = kv.size(2); + const int64_t emb_dim = k_rope.size(-1); + TORCH_CHECK(k_rope.dim() >= 3 && k_rope.size(0) == s && k_rope.size(1) == b && + k_rope.numel() == s * b * emb_dim, + "expected k_rope of shape [s, b, 1, emb_dim] or [s, b, emb_dim], got ", + k_rope.sizes()); + TORCH_CHECK(v_head_dim > 0 && v_head_dim < kv.size(3), "invalid v_head_dim ", v_head_dim, + " for kv of shape ", kv.sizes()); + const int64_t key_dim = kv.size(3) - v_head_dim + emb_dim; + // k_rope rows are read in place when they are a fixed, 16-byte aligned stride apart. + at::Tensor k_rope_rows = k_rope; + int64_t k_rope_stride = b > 1 ? k_rope_rows.stride(1) : k_rope_rows.stride(0); + if (k_rope_rows.stride(-1) != 1 || + (b > 1 && k_rope_rows.stride(0) != b * k_rope_rows.stride(1)) || k_rope_stride < emb_dim || + k_rope_stride % 8 != 0 || reinterpret_cast(k_rope_rows.data_ptr()) % 16 != 0) { + k_rope_rows = k_rope.contiguous(); + k_rope_stride = emb_dim; + } + const auto kv_c = kv.contiguous(); + const auto cos_c = cos.contiguous(); + const auto sin_c = sin.contiguous(); + auto key = empty_mxfp8_attn_buffers(s, b, h, key_dim, kv.device()); + auto value = empty_mxfp8_attn_buffers(s, b, h, v_head_dim, kv.device()); + auto kv_cu = makeTransformerEngineTensor(kv_c); + auto k_rope_cu = makeTransformerEngineTensor( + k_rope_rows.data_ptr(), + std::vector{static_cast(s * b), static_cast(emb_dim)}, + GetTransformerEngineDType(k_rope_rows.scalar_type())); + auto cos_cu = makeTransformerEngineTensor(cos_c); + auto sin_cu = makeTransformerEngineTensor(sin_c); + auto key_cu = make_swizzled_mxfp8_tensor(key); + auto value_cu = make_swizzled_mxfp8_tensor(value); + NVTE_SCOPED_GIL_RELEASE({ + nvte_fused_mla_kv_rope_mxfp8_attn( + kv_cu.data(), k_rope_cu.data(), static_cast(k_rope_stride), cos_cu.data(), + sin_cu.data(), key_cu.data(), value_cu.data(), at::cuda::getCurrentCUDAStream()); + }); + return {key, value}; +} + +at::Tensor fused_mla_q_rope_backward_(at::Tensor dq, const at::Tensor &cos, const at::Tensor &sin) { + check_mla_rope_inputs(dq, "dq", {cos, sin}); + TORCH_CHECK(dq.is_contiguous(), "expected a contiguous dq"); + at::cuda::CUDAGuard device_guard(dq.device()); + const auto cos_c = cos.contiguous(); + const auto sin_c = sin.contiguous(); + auto dq_cu = makeTransformerEngineTensor(dq); + auto cos_cu = makeTransformerEngineTensor(cos_c); + auto sin_cu = makeTransformerEngineTensor(sin_c); + NVTE_SCOPED_GIL_RELEASE({ + nvte_fused_mla_q_rope_backward(dq_cu.data(), cos_cu.data(), sin_cu.data(), + at::cuda::getCurrentCUDAStream()); + }); + return dq; +} + +std::vector fused_mla_q_rope_backward_mxfp8(const at::Tensor &dq, const at::Tensor &cos, + const at::Tensor &sin) { + check_mla_rope_inputs(dq, "dq", {cos, sin}); + at::cuda::CUDAGuard device_guard(dq.device()); + const auto dq_c = dq.contiguous(); + const auto cos_c = cos.contiguous(); + const auto sin_c = sin.contiguous(); + const int64_t tokens = dq.size(0) * dq.size(1), columns = dq.size(2) * dq.size(3); + auto out = empty_mxfp8_buffers({tokens, columns}, {tokens, columns / 32}, {tokens / 32, columns}, + dq.device()); + auto dq_cu = makeTransformerEngineTensor(dq_c); + auto cos_cu = makeTransformerEngineTensor(cos_c); + auto sin_cu = makeTransformerEngineTensor(sin_c); + auto out_cu = make_swizzled_mxfp8_tensor(out); + NVTE_SCOPED_GIL_RELEASE({ + nvte_fused_mla_q_rope_backward_mxfp8(dq_cu.data(), cos_cu.data(), sin_cu.data(), out_cu.data(), + at::cuda::getCurrentCUDAStream()); + }); + return out; +} + +std::vector fused_mla_kv_split_backward_mxfp8(const at::Tensor &dk, + const at::Tensor &dv, + const int64_t k_dim) { + check_mla_rope_inputs(dk, "dk", {dv}); + TORCH_CHECK(dv.dim() == 4, "expected dv of shape [s, b, h, d], got ", dv.sizes()); + at::cuda::CUDAGuard device_guard(dk.device()); + const auto dk_c = dk.contiguous(); + const auto dv_c = dv.contiguous(); + const int64_t tokens = dk.size(0) * dk.size(1); + TORCH_CHECK(k_dim > 0 && k_dim < dk.size(3), "invalid k_dim ", k_dim, " for dk of shape ", + dk.sizes()); + const int64_t columns = dk.size(2) * (k_dim + dv.size(3)); + auto out = empty_mxfp8_buffers({tokens, columns}, {tokens, columns / 32}, {tokens / 32, columns}, + dk.device()); + auto dk_cu = makeTransformerEngineTensor(dk_c); + auto dv_cu = makeTransformerEngineTensor(dv_c); + auto out_cu = make_swizzled_mxfp8_tensor(out); + NVTE_SCOPED_GIL_RELEASE({ + nvte_fused_mla_kv_split_backward_mxfp8(dk_cu.data(), dv_cu.data(), out_cu.data(), + at::cuda::getCurrentCUDAStream()); + }); + return out; +} + } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 6fbed1436fb..b8d33430a9a 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -613,6 +613,23 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Split MLA key/value, apply RoPE to the key and quantize both to MXFP8", py::arg("kv"), py::arg("k_rope"), py::arg("cos"), py::arg("sin"), py::arg("v_head_dim"), py::arg("key_quantizer"), py::arg("value_quantizer")); + m.def("fused_mla_q_rope_mxfp8_attn", &transformer_engine::pytorch::fused_mla_q_rope_mxfp8_attn, + "Apply RoPE to the MLA query and quantize it to MXFP8 with the attention scale layout", + py::arg("q"), py::arg("cos"), py::arg("sin")); + m.def("fused_mla_kv_rope_mxfp8_attn", &transformer_engine::pytorch::fused_mla_kv_rope_mxfp8_attn, + "Split MLA key/value, apply RoPE to the key and quantize both to MXFP8 with the attention " + "scale layout", + py::arg("kv"), py::arg("k_rope"), py::arg("cos"), py::arg("sin"), py::arg("v_head_dim")); + m.def("fused_mla_q_rope_backward_", &transformer_engine::pytorch::fused_mla_q_rope_backward_, + "MLA query RoPE backward, in place", py::arg("dq"), py::arg("cos"), py::arg("sin")); + m.def("fused_mla_q_rope_backward_mxfp8", + &transformer_engine::pytorch::fused_mla_q_rope_backward_mxfp8, + "MLA query RoPE backward quantized to MXFP8 (GEMM-swizzled scales)", py::arg("dq"), + py::arg("cos"), py::arg("sin")); + m.def("fused_mla_kv_split_backward_mxfp8", + &transformer_engine::pytorch::fused_mla_kv_split_backward_mxfp8, + "MLA key/value split backward quantized to MXFP8 (GEMM-swizzled scales)", py::arg("dk"), + py::arg("dv"), py::arg("k_dim")); // fused router transformer_engine::pytorch::init_router_bindings(m); From 5daf5f56ac943891a4c57e81e489789a191f95a8 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman S Date: Fri, 13 Nov 2026 11:03:24 -0800 Subject: [PATCH 2/2] [PyTorch] Test MLA RoPE MXFP8 kernels with asymmetric RoPE tables build_rope_tables repeats the same frequencies in both halves, so the q/kv RoPE forward and the q RoPE backward tests could not catch a kernel that reads the wrong half of cos / sin. Run them with an asymmetric table as well. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Siddhartha Raman S --- tests/pytorch/attention/test_fused_mla.py | 44 ++++++++++++++++++----- 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/tests/pytorch/attention/test_fused_mla.py b/tests/pytorch/attention/test_fused_mla.py index bc10a0bf0a4..4d28a70ce87 100644 --- a/tests/pytorch/attention/test_fused_mla.py +++ b/tests/pytorch/attention/test_fused_mla.py @@ -479,11 +479,31 @@ def test_inputs_on_non_current_device(): # --- MXFP8 attention scale layout and the RoPE backward kernels ----------------------------------- -def _rope_tables(s, dtype=torch.bfloat16, rows=None): - cos, sin = build_rope_tables(rows or s, emb_dim=HEAD_DIM_ROPE, device="cuda") +def _rope_tables(s, dtype=torch.bfloat16, rows=None, asymmetric=False): + """RoPE tables scaled by 1.2 (as in YaRN, so the products are not exact). build_rope_tables + repeats the same frequencies in both halves; with asymmetric=True the right half uses + different frequencies, so a kernel that reads the wrong half of cos / sin fails.""" + rows = rows or s + if not asymmetric: + cos, sin = build_rope_tables(rows, emb_dim=HEAD_DIM_ROPE, device="cuda") + else: + inv_freq = 1.0 / ( + 10000 + ** ( + torch.arange(0, HEAD_DIM_ROPE, 2, dtype=torch.float32, device="cuda") + / HEAD_DIM_ROPE + ) + ) + t = torch.arange(rows, device="cuda", dtype=torch.float32) + freqs = torch.cat([torch.outer(t, inv_freq), torch.outer(t, inv_freq * 0.5)], dim=-1) + cos, sin = freqs.cos(), freqs.sin() + assert not torch.equal(cos[1:, : HEAD_DIM_ROPE // 2], cos[1:, HEAD_DIM_ROPE // 2 :]) return (cos * 1.2).to(dtype), (sin * 1.2).to(dtype) +TABLES = pytest.mark.parametrize("asymmetric", [False, True], ids=["symmetric", "asymmetric"]) + + def _q_input(s, b, h, seed=4321): torch.manual_seed(seed) q = (torch.randn(s, b, h, HEAD_DIM, device="cuda") * 4).to(torch.bfloat16) @@ -515,10 +535,11 @@ def _assert_swizzled_identical(actual, expected, name): @requires_kv_rope_kernel +@TABLES @pytest.mark.parametrize("s, b, h", ATTN_SHAPES) -def test_q_rope_mxfp8_matches_rope_quantize_swizzle(s, b, h): +def test_q_rope_mxfp8_matches_rope_quantize_swizzle(s, b, h, asymmetric): q = _q_input(s, b, h) - cos, sin = _rope_tables(s) + cos, sin = _rope_tables(s, asymmetric=asymmetric) expected = _q_reference(q, cos, sin) q_before = q.clone() query = fused_mla_q_rope_mxfp8(q, cos.view(s, 1, 1, -1), sin.view(s, 1, 1, -1)) @@ -530,9 +551,12 @@ def test_q_rope_mxfp8_matches_rope_quantize_swizzle(s, b, h): @requires_kv_rope_kernel +@TABLES @pytest.mark.parametrize("s, b, h", ATTN_SHAPES) -def test_kv_rope_mxfp8_transpose_swizzle_scales(s, b, h): +def test_kv_rope_mxfp8_transpose_swizzle_scales(s, b, h, asymmetric): kv, k_rope, cos, sin = _kv_rope_inputs(s, b, h, torch.bfloat16) + if asymmetric: + cos, sin = _rope_tables(s, asymmetric=True) ref_key, ref_value = _kv_rope_reference(kv, k_rope, cos, sin) mxfp8_transpose_swizzle([ref_key, ref_value], "sbhd") key, value = fused_mla_kv_rope_mxfp8( @@ -630,10 +654,11 @@ def _grad_inputs(s, b, h, seed=99): @requires_kv_rope_kernel +@TABLES @pytest.mark.parametrize("s, b, h", ATTN_SHAPES + [(64, 2, 4)]) -def test_q_rope_backward_in_place(s, b, h): +def test_q_rope_backward_in_place(s, b, h, asymmetric): dq, _, _ = _grad_inputs(s, b, h) - cos, sin = _rope_tables(s) + cos, sin = _rope_tables(s, asymmetric=asymmetric) expected = _q_rope_backward_reference(dq, cos, sin) out = fused_mla_q_rope_backward_(dq, cos, sin) assert out.data_ptr() == dq.data_ptr() @@ -641,10 +666,11 @@ def test_q_rope_backward_in_place(s, b, h): @requires_kv_rope_kernel +@TABLES @pytest.mark.parametrize("s, b, h", ATTN_SHAPES + [(64, 2, 4)]) -def test_q_rope_backward_mxfp8(s, b, h): +def test_q_rope_backward_mxfp8(s, b, h, asymmetric): dq, _, _ = _grad_inputs(s, b, h) - cos, sin = _rope_tables(s) + cos, sin = _rope_tables(s, asymmetric=asymmetric) expected_bf16 = _q_rope_backward_reference(dq, cos, sin) expected = _grad_quantizer()(expected_bf16.view(s * b, h * HEAD_DIM)) dq_before = dq.clone()