diff --git a/transformer_engine/plugin/core/backends/flagos/flagos.py b/transformer_engine/plugin/core/backends/flagos/flagos.py index f651be22e0..21e065ce39 100644 --- a/transformer_engine/plugin/core/backends/flagos/flagos.py +++ b/transformer_engine/plugin/core/backends/flagos/flagos.py @@ -22,6 +22,10 @@ scaled_masked_softmax_forward_fl, scaled_masked_softmax_backward_fl, te_general_grouped_gemm_fl, + fused_rope_forward_fl, + fused_rope_backward_fl, + fused_qkv_rope_forward_fl, + fused_qkv_rope_backward_fl, ) @@ -352,6 +356,101 @@ def multi_tensor_adam_param_remainder( weight_decay, ) + # fused apply rope + def fused_rope_forward( + self, + input: torch.Tensor, + freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: + return fused_rope_forward_fl( + input, + freqs, + start_positions, + qkv_format, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ) + + def fused_rope_backward( + self, + output_grads: torch.Tensor, + freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: + return fused_rope_backward_fl( + output_grads, + freqs, + start_positions, + qkv_format, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ) + + def fused_qkv_rope_forward( + self, + qkv_input: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_split_arg_list: List[int], + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cp_size: int, + cp_rank: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + return fused_qkv_rope_forward_fl( + qkv_input, + q_freqs, + k_freqs, + start_positions, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + + def fused_qkv_rope_backward( + self, + q_grad_out: torch.Tensor, + k_grad_out: torch.Tensor, + v_grad_out: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + qkv_split_arg_list: List[int], + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: + return fused_qkv_rope_backward_fl( + q_grad_out, + k_grad_out, + v_grad_out, + q_freqs, + k_freqs, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + # Misc def get_cublasLt_version(self) -> int: return 110000 diff --git a/transformer_engine/plugin/core/backends/flagos/impl/__init__.py b/transformer_engine/plugin/core/backends/flagos/impl/__init__.py index db0381f259..f270ffef3b 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/__init__.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/__init__.py @@ -8,3 +8,4 @@ from .multi_tensor import * from .softmax import * from .normalization import * +from .trition import * diff --git a/transformer_engine/plugin/core/backends/flagos/impl/trition/__init__.py b/transformer_engine/plugin/core/backends/flagos/impl/trition/__init__.py new file mode 100644 index 0000000000..e1bbcbdb0e --- /dev/null +++ b/transformer_engine/plugin/core/backends/flagos/impl/trition/__init__.py @@ -0,0 +1,5 @@ +# Copyright (c) 2025, BAAI. All rights reserved. +# +# See LICENSE for license information. + +from .fused_rope import * diff --git a/transformer_engine/plugin/core/backends/flagos/impl/trition/fused_rope.py b/transformer_engine/plugin/core/backends/flagos/impl/trition/fused_rope.py new file mode 100644 index 0000000000..9840856ee6 --- /dev/null +++ b/transformer_engine/plugin/core/backends/flagos/impl/trition/fused_rope.py @@ -0,0 +1,750 @@ +# Copyright (c) 2025, BAAI. All rights reserved. +# +# See LICENSE for license information. + +from __future__ import annotations + +from typing import List, Optional, Tuple + +import torch + +try: + import triton + import triton.language as tl +except ModuleNotFoundError: # pragma: no cover - exercised only on systems without Triton. + triton = None + tl = None + + +NVTE_SBHD = 0 +NVTE_BSHD = 1 +NVTE_THD = 2 + +__all__ = [ + "fused_rope_forward_fl", + "fused_rope_backward_fl", + "fused_qkv_rope_forward_fl", + "fused_qkv_rope_backward_fl", +] + + +def _require_triton() -> None: + if triton is None: + raise RuntimeError( + "FlagOS fused RoPE requires the Triton Python package, but it is not installed." + ) + + +def _next_power_of_2(value: int) -> int: + return 1 << (value - 1).bit_length() + + +def _choose_block_d(d: int) -> int: + return min(max(16, _next_power_of_2(min(d, 128))), 128) + + +def _choose_rope_block_h(h: int) -> int: + return 4 if h < 16 else 8 + + +def _choose_qkv_block_h(h: int) -> int: + return min(8, _next_power_of_2(max(1, h))) + + +def _num_warps(block_h: int) -> int: + return max(1, min(8, block_h)) + + +def _check_freqs(freqs: torch.Tensor, name: str) -> None: + if freqs.dim() != 4: + raise ValueError(f"{name} must be a 4D tensor") + if freqs.size(1) != 1 or freqs.size(2) != 1: + raise ValueError(f"{name} must have shape (s, 1, 1, d)") + if freqs.dtype != torch.float32: + raise TypeError(f"{name} must have dtype torch.float32") + + +def _check_qkv_splits(qkv_split_arg_list: List[int]) -> Tuple[int, int, int]: + if len(qkv_split_arg_list) != 3: + raise ValueError("qkv_split_arg_list must contain exactly three integers") + q_split, k_split, v_split = [int(x) for x in qkv_split_arg_list] + if q_split <= 0 or k_split <= 0 or v_split <= 0: + raise ValueError("qkv split sizes must be positive") + if k_split != v_split: + raise ValueError("FlagOS fused QKV RoPE requires equal K and V head dimensions") + if q_split % k_split != 0: + raise ValueError("Q split size must be an integer multiple of the K/V head dimension") + return q_split, k_split, v_split + + +if triton is not None: + + @triton.jit + def _fused_rope_kernel( + src, + cu_seqlens, + freqs, + start_positions, + dst, + S: tl.constexpr, + B: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + D2: tl.constexpr, + STRIDE_S_OR_T: tl.constexpr, + STRIDE_B: tl.constexpr, + STRIDE_H: tl.constexpr, + STRIDE_D: tl.constexpr, + QKV_FORMAT: tl.constexpr, + INTERLEAVED: tl.constexpr, + IS_BACKWARD: tl.constexpr, + HAS_CU_SEQLENS: tl.constexpr, + HAS_START_POSITIONS: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_RANK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + N_D_BLOCKS: tl.constexpr, + ): + s_id = tl.program_id(0) + b_id = tl.program_id(1) + hd_pid = tl.program_id(2) + h_block = hd_pid // N_D_BLOCKS + d_block = hd_pid - h_block * N_D_BLOCKS + + offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D) + mask_h = offs_h < H + mask_d = offs_d < D + mask_d2 = offs_d < D2 + mask = mask_h[:, None] & mask_d[None, :] + mask_rotary = mask_h[:, None] & mask_d2[None, :] + + if HAS_CU_SEQLENS: + start = tl.load(cu_seqlens + b_id) // CP_SIZE + end = tl.load(cu_seqlens + b_id + 1) // CP_SIZE + t_id = s_id + start + valid_token = t_id < end + offset_block = t_id * STRIDE_S_OR_T + offset_block_dst = t_id * H * D + cur_seqlens = end - start + else: + valid_token = True + offset_block = s_id * STRIDE_S_OR_T + b_id * STRIDE_B + if QKV_FORMAT == 0: + offset_block_dst = s_id * B * H * D + b_id * H * D + else: + offset_block_dst = b_id * S * H * D + s_id * H * D + cur_seqlens = S + + begin_offset = 0 + if HAS_START_POSITIONS: + begin_offset = tl.load(start_positions + b_id) + s_id_for_freqs = s_id + begin_offset + + if CP_SIZE > 1: + half_seq = cur_seqlens // 2 + cp_delta = tl.where( + s_id < half_seq, + CP_RANK * half_seq, + cur_seqlens * CP_SIZE - (CP_RANK + 1) * half_seq - half_seq, + ) + s_id_for_freqs += cp_delta + + src_offsets = offset_block + offs_h[:, None] * STRIDE_H + offs_d[None, :] * STRIDE_D + dst_offsets = offset_block_dst + offs_h[:, None] * D + offs_d[None, :] + + src_values = tl.load(src + src_offsets, mask=mask & valid_token, other=0.0).to(tl.float32) + out_values = src_values + + if INTERLEAVED: + is_even = (offs_d % 2) == 0 + if IS_BACKWARD: + rot_d = tl.where(is_even, offs_d + 1, offs_d - 1) + sin_d = rot_d + sin_sign = tl.where(is_even, 1.0, -1.0) + rot_sign = 1.0 + else: + rot_d = tl.where(is_even, offs_d + 1, offs_d - 1) + sin_d = offs_d + sin_sign = 1.0 + rot_sign = tl.where(is_even, -1.0, 1.0) + else: + half_d2 = D2 // 2 + first_half = (offs_d + half_d2) < D2 + rot_d = tl.where(first_half, offs_d + half_d2, offs_d + half_d2 - D2) + if IS_BACKWARD: + sin_d = rot_d + sin_sign = tl.where(first_half, 1.0, -1.0) + rot_sign = 1.0 + else: + sin_d = offs_d + sin_sign = 1.0 + rot_sign = tl.where(first_half, -1.0, 1.0) + + rot_offsets = offset_block + offs_h[:, None] * STRIDE_H + rot_d[None, :] * STRIDE_D + rot_values = tl.load(src + rot_offsets, mask=mask_rotary & valid_token, other=0.0).to( + tl.float32 + ) + freq_base = s_id_for_freqs * D2 + freq_mask = mask_d2 & valid_token + cos_values = tl.cos(tl.load(freqs + freq_base + offs_d, mask=freq_mask, other=0.0)) + sin_values = ( + tl.sin(tl.load(freqs + freq_base + sin_d, mask=freq_mask, other=0.0)) * sin_sign + ) + rotary_values = ( + src_values * cos_values[None, :] + rot_values * rot_sign * sin_values[None, :] + ) + out_values = tl.where(mask_d2[None, :], rotary_values, out_values) + + tl.store(dst + dst_offsets, out_values, mask=mask & valid_token) + + @triton.jit + def _fused_qkv_rope_kernel( + qkv_input, + q_freqs, + k_freqs, + start_positions, + q_out, + k_out, + v_out, + qkv_grad_input, + S: tl.constexpr, + B: tl.constexpr, + H: tl.constexpr, + D: tl.constexpr, + D2: tl.constexpr, + Q_SPLIT: tl.constexpr, + K_SPLIT: tl.constexpr, + V_SPLIT: tl.constexpr, + QKV_FORMAT: tl.constexpr, + INTERLEAVED: tl.constexpr, + IS_BACKWARD: tl.constexpr, + HAS_START_POSITIONS: tl.constexpr, + CP_SIZE: tl.constexpr, + CP_RANK: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + N_D_BLOCKS: tl.constexpr, + ): + s_id = tl.program_id(0) + b_id = tl.program_id(1) + hd_pid = tl.program_id(2) + h_block = hd_pid // N_D_BLOCKS + d_block = hd_pid - h_block * N_D_BLOCKS + + offs_h = h_block * BLOCK_H + tl.arange(0, BLOCK_H) + offs_d = d_block * BLOCK_D + tl.arange(0, BLOCK_D) + mask_h = offs_h < H + mask_d = offs_d < D + mask_d2 = offs_d < D2 + + total_d = Q_SPLIT + K_SPLIT + V_SPLIT + if QKV_FORMAT == 0: + input_base = s_id * B * H * total_d + b_id * H * total_d + q_base = s_id * B * H * Q_SPLIT + b_id * H * Q_SPLIT + k_base = s_id * B * H * K_SPLIT + b_id * H * K_SPLIT + v_base = s_id * B * H * V_SPLIT + b_id * H * V_SPLIT + else: + input_base = b_id * S * H * total_d + s_id * H * total_d + q_base = b_id * S * H * Q_SPLIT + s_id * H * Q_SPLIT + k_base = b_id * S * H * K_SPLIT + s_id * H * K_SPLIT + v_base = b_id * S * H * V_SPLIT + s_id * H * V_SPLIT + + if CP_SIZE > 1: + half_seq = S // 2 + s_id_for_freqs = tl.where( + s_id < half_seq, + s_id + CP_RANK * half_seq, + S * CP_SIZE - (CP_RANK + 1) * half_seq + s_id - half_seq, + ) + else: + if IS_BACKWARD: + s_id_for_freqs = s_id + else: + begin_offset = 0 + if HAS_START_POSITIONS: + begin_offset = tl.load(start_positions + b_id) + s_id_for_freqs = s_id + begin_offset + + if INTERLEAVED: + is_even = (offs_d % 2) == 0 + if IS_BACKWARD: + rot_d = tl.where(is_even, offs_d + 1, offs_d - 1) + sin_d = rot_d + sin_sign = tl.where(is_even, 1.0, -1.0) + rot_sign = 1.0 + else: + rot_d = tl.where(is_even, offs_d + 1, offs_d - 1) + sin_d = offs_d + sin_sign = 1.0 + rot_sign = tl.where(is_even, -1.0, 1.0) + else: + half_d2 = D2 // 2 + first_half = (offs_d + half_d2) < D2 + rot_d = tl.where(first_half, offs_d + half_d2, offs_d + half_d2 - D2) + if IS_BACKWARD: + sin_d = rot_d + sin_sign = tl.where(first_half, 1.0, -1.0) + rot_sign = 1.0 + else: + sin_d = offs_d + sin_sign = 1.0 + rot_sign = tl.where(first_half, -1.0, 1.0) + + q_cos = tl.cos(tl.load(q_freqs + s_id_for_freqs * D2 + offs_d, mask=mask_d2, other=0.0)) + q_sin = ( + tl.sin(tl.load(q_freqs + s_id_for_freqs * D2 + sin_d, mask=mask_d2, other=0.0)) + * sin_sign + ) + k_cos = tl.cos(tl.load(k_freqs + s_id_for_freqs * D2 + offs_d, mask=mask_d2, other=0.0)) + k_sin = ( + tl.sin(tl.load(k_freqs + s_id_for_freqs * D2 + sin_d, mask=mask_d2, other=0.0)) + * sin_sign + ) + + for row_offset in tl.static_range(0, Q_SPLIT, D): + component_d = row_offset + offs_d + mask = mask_h[:, None] & (component_d[None, :] < Q_SPLIT) & mask_d[None, :] + mask_rotary = mask_h[:, None] & (component_d[None, :] < Q_SPLIT) & mask_d2[None, :] + if IS_BACKWARD: + src_base = q_base + dst_base = input_base + src_row_length = Q_SPLIT + dst_row_offset = row_offset + else: + src_base = input_base + dst_base = q_base + src_row_length = total_d + dst_row_offset = row_offset + src_offsets = src_base + offs_h[:, None] * src_row_length + component_d[None, :] + rot_offsets = ( + src_base + offs_h[:, None] * src_row_length + (row_offset + rot_d)[None, :] + ) + dst_offsets = dst_base + offs_h[:, None] * total_d + dst_row_offset + offs_d[None, :] + if not IS_BACKWARD: + dst_offsets = dst_base + offs_h[:, None] * Q_SPLIT + component_d[None, :] + + if IS_BACKWARD: + values = tl.load(q_out + src_offsets, mask=mask, other=0.0).to(tl.float32) + rot_values = tl.load(q_out + rot_offsets, mask=mask_rotary, other=0.0).to( + tl.float32 + ) + else: + values = tl.load(qkv_input + src_offsets, mask=mask, other=0.0).to(tl.float32) + rot_values = tl.load(qkv_input + rot_offsets, mask=mask_rotary, other=0.0).to( + tl.float32 + ) + rotary_values = values * q_cos[None, :] + rot_values * rot_sign * q_sin[None, :] + out_values = tl.where(mask_d2[None, :], rotary_values, values) + if IS_BACKWARD: + tl.store(qkv_grad_input + dst_offsets, out_values, mask=mask) + else: + tl.store(q_out + dst_offsets, out_values, mask=mask) + + for row_offset in tl.static_range(0, K_SPLIT, D): + component_d = row_offset + offs_d + input_row_offset = Q_SPLIT + row_offset + mask = mask_h[:, None] & (component_d[None, :] < K_SPLIT) & mask_d[None, :] + mask_rotary = mask_h[:, None] & (component_d[None, :] < K_SPLIT) & mask_d2[None, :] + if IS_BACKWARD: + src_offsets = k_base + offs_h[:, None] * K_SPLIT + component_d[None, :] + rot_offsets = k_base + offs_h[:, None] * K_SPLIT + (row_offset + rot_d)[None, :] + dst_offsets = ( + input_base + offs_h[:, None] * total_d + input_row_offset + offs_d[None, :] + ) + values = tl.load(k_out + src_offsets, mask=mask, other=0.0).to(tl.float32) + rot_values = tl.load(k_out + rot_offsets, mask=mask_rotary, other=0.0).to( + tl.float32 + ) + rotary_values = values * k_cos[None, :] + rot_values * rot_sign * k_sin[None, :] + out_values = tl.where(mask_d2[None, :], rotary_values, values) + tl.store(qkv_grad_input + dst_offsets, out_values, mask=mask) + else: + src_offsets = ( + input_base + offs_h[:, None] * total_d + input_row_offset + offs_d[None, :] + ) + rot_offsets = ( + input_base + offs_h[:, None] * total_d + (input_row_offset + rot_d)[None, :] + ) + dst_offsets = k_base + offs_h[:, None] * K_SPLIT + component_d[None, :] + values = tl.load(qkv_input + src_offsets, mask=mask, other=0.0).to(tl.float32) + rot_values = tl.load(qkv_input + rot_offsets, mask=mask_rotary, other=0.0).to( + tl.float32 + ) + rotary_values = values * k_cos[None, :] + rot_values * rot_sign * k_sin[None, :] + out_values = tl.where(mask_d2[None, :], rotary_values, values) + tl.store(k_out + dst_offsets, out_values, mask=mask) + + component_d = offs_d + mask = mask_h[:, None] & (component_d[None, :] < V_SPLIT) & mask_d[None, :] + if IS_BACKWARD: + src_offsets = v_base + offs_h[:, None] * V_SPLIT + component_d[None, :] + dst_offsets = ( + input_base + offs_h[:, None] * total_d + Q_SPLIT + K_SPLIT + offs_d[None, :] + ) + values = tl.load(v_out + src_offsets, mask=mask, other=0.0) + tl.store(qkv_grad_input + dst_offsets, values, mask=mask) + else: + src_offsets = ( + input_base + offs_h[:, None] * total_d + Q_SPLIT + K_SPLIT + offs_d[None, :] + ) + dst_offsets = v_base + offs_h[:, None] * V_SPLIT + component_d[None, :] + values = tl.load(qkv_input + src_offsets, mask=mask, other=0.0) + tl.store(v_out + dst_offsets, values, mask=mask) + + +def fused_rope_forward_fl( + input: torch.Tensor, + freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, +) -> torch.Tensor: + _require_triton() + _check_freqs(freqs, "freqs") + if not freqs.is_contiguous(): + freqs = freqs.contiguous() + qkv_format = int(qkv_format) + output = torch.empty(input.size(), dtype=input.dtype, device=input.device) + + if qkv_format == NVTE_THD: + if input.dim() != 3: + raise ValueError("input must be a 3D tensor for THD format") + if cu_seqlens is None: + raise ValueError("cu_seqlens is required for THD format") + s = freqs.size(0) + b = cu_seqlens.numel() - 1 + h = input.size(1) + d = input.size(2) + stride_s_or_t = input.stride(0) + stride_b = 0 + stride_h = input.stride(1) + stride_d = input.stride(2) + has_cu_seqlens = True + else: + if input.dim() != 4: + raise ValueError("input must be a 4D tensor for SBHD/BSHD format") + if qkv_format == NVTE_SBHD: + s = input.size(0) + b = input.size(1) + stride_s_or_t = input.stride(0) + stride_b = input.stride(1) + else: + s = input.size(1) + b = input.size(0) + stride_s_or_t = input.stride(1) + stride_b = input.stride(0) + h = input.size(2) + d = input.size(3) + stride_h = input.stride(2) + stride_d = input.stride(3) + has_cu_seqlens = False + + d2 = freqs.size(3) + if d < d2: + raise ValueError("input last dimension must be greater than or equal to freqs last dim") + if qkv_format != NVTE_THD and s * cp_size > freqs.size(0): + raise ValueError("freqs sequence length is too short for input and cp_size") + + block_h = _choose_rope_block_h(h) + block_d = _choose_block_d(d) + d_blocks = triton.cdiv(d, block_d) + grid = (s, b, triton.cdiv(h, block_h) * d_blocks) + dummy_cu = cu_seqlens if cu_seqlens is not None else input + dummy_start = start_positions if start_positions is not None else input + _fused_rope_kernel[grid]( + input, + dummy_cu, + freqs, + dummy_start, + output, + s, + b, + h, + d, + d2, + stride_s_or_t, + stride_b, + stride_h, + stride_d, + qkv_format, + interleaved, + False, + has_cu_seqlens, + start_positions is not None, + cp_size, + cp_rank, + BLOCK_H=block_h, + BLOCK_D=block_d, + N_D_BLOCKS=d_blocks, + num_warps=_num_warps(block_h), + ) + return output + + +def fused_rope_backward_fl( + output_grads: torch.Tensor, + freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, +) -> torch.Tensor: + _require_triton() + _check_freqs(freqs, "freqs") + if not freqs.is_contiguous(): + freqs = freqs.contiguous() + qkv_format = int(qkv_format) + input_grads = torch.empty( + output_grads.size(), dtype=output_grads.dtype, device=output_grads.device + ) + + if qkv_format == NVTE_THD: + if output_grads.dim() != 3: + raise ValueError("output_grads must be a 3D tensor for THD format") + if cu_seqlens is None: + raise ValueError("cu_seqlens is required for THD format") + s = freqs.size(0) + b = cu_seqlens.numel() - 1 + h = output_grads.size(1) + d = output_grads.size(2) + stride_s_or_t = output_grads.stride(0) + stride_b = 0 + stride_h = output_grads.stride(1) + stride_d = output_grads.stride(2) + has_cu_seqlens = True + else: + if output_grads.dim() != 4: + raise ValueError("output_grads must be a 4D tensor for SBHD/BSHD format") + if qkv_format == NVTE_SBHD: + s = output_grads.size(0) + b = output_grads.size(1) + stride_s_or_t = output_grads.stride(0) + stride_b = output_grads.stride(1) + else: + s = output_grads.size(1) + b = output_grads.size(0) + stride_s_or_t = output_grads.stride(1) + stride_b = output_grads.stride(0) + h = output_grads.size(2) + d = output_grads.size(3) + stride_h = output_grads.stride(2) + stride_d = output_grads.stride(3) + has_cu_seqlens = False + + d2 = freqs.size(3) + if d < d2: + raise ValueError( + "output_grads last dimension must be greater than or equal to freqs last dim" + ) + if qkv_format != NVTE_THD and s * cp_size > freqs.size(0): + raise ValueError("freqs sequence length is too short for output_grads and cp_size") + + block_h = _choose_rope_block_h(h) + block_d = _choose_block_d(d) + d_blocks = triton.cdiv(d, block_d) + grid = (s, b, triton.cdiv(h, block_h) * d_blocks) + dummy_cu = cu_seqlens if cu_seqlens is not None else output_grads + dummy_start = start_positions if start_positions is not None else output_grads + _fused_rope_kernel[grid]( + output_grads, + dummy_cu, + freqs, + dummy_start, + input_grads, + s, + b, + h, + d, + d2, + stride_s_or_t, + stride_b, + stride_h, + stride_d, + qkv_format, + interleaved, + True, + has_cu_seqlens, + start_positions is not None, + cp_size, + cp_rank, + BLOCK_H=block_h, + BLOCK_D=block_d, + N_D_BLOCKS=d_blocks, + num_warps=_num_warps(block_h), + ) + return input_grads + + +def fused_qkv_rope_forward_fl( + qkv_input: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_split_arg_list: List[int], + qkv_format, + interleaved: bool, + cp_size: int, + cp_rank: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + _require_triton() + _check_freqs(q_freqs, "q_freqs") + _check_freqs(k_freqs, "k_freqs") + if not q_freqs.is_contiguous(): + q_freqs = q_freqs.contiguous() + if not k_freqs.is_contiguous(): + k_freqs = k_freqs.contiguous() + if qkv_input.dim() != 4: + raise ValueError("qkv_input must be a 4D tensor") + if not qkv_input.is_contiguous(): + raise ValueError("qkv_input must be contiguous") + + qkv_format = int(qkv_format) + is_sbhd = qkv_format == NVTE_SBHD + s = qkv_input.size(0) if is_sbhd else qkv_input.size(1) + b = qkv_input.size(1) if is_sbhd else qkv_input.size(0) + h = qkv_input.size(2) + q_split, k_split, v_split = _check_qkv_splits(qkv_split_arg_list) + if qkv_input.size(3) != q_split + k_split + v_split: + raise ValueError("qkv_input last dimension must equal the sum of qkv split sizes") + d = v_split + d2 = q_freqs.size(3) + if d < d2: + raise ValueError("qkv value split must be greater than or equal to q_freqs last dim") + if q_freqs.size(3) != k_freqs.size(3): + raise ValueError("q_freqs and k_freqs must have the same rotary dimension") + + q_out_size = list(qkv_input.size()) + q_out_size[2] = q_out_size[2] * q_split // k_split + q_out_size[3] = k_split + k_out_size = list(qkv_input.size()) + k_out_size[3] = k_split + v_out_size = list(qkv_input.size()) + v_out_size[3] = v_split + q_out = torch.empty(q_out_size, dtype=qkv_input.dtype, device=qkv_input.device) + k_out = torch.empty(k_out_size, dtype=qkv_input.dtype, device=qkv_input.device) + v_out = torch.empty(v_out_size, dtype=qkv_input.dtype, device=qkv_input.device) + + block_h = _choose_qkv_block_h(h) + block_d = _choose_block_d(d) + d_blocks = triton.cdiv(d, block_d) + grid = (s, b, triton.cdiv(h, block_h) * d_blocks) + dummy_start = start_positions if start_positions is not None else qkv_input + _fused_qkv_rope_kernel[grid]( + qkv_input, + q_freqs, + k_freqs, + dummy_start, + q_out, + k_out, + v_out, + qkv_input, + s, + b, + h, + d, + d2, + q_split, + k_split, + v_split, + qkv_format, + interleaved, + False, + start_positions is not None, + cp_size, + cp_rank, + BLOCK_H=block_h, + BLOCK_D=block_d, + N_D_BLOCKS=d_blocks, + num_warps=_num_warps(block_h), + ) + return q_out, k_out, v_out + + +def fused_qkv_rope_backward_fl( + q_grad_out: torch.Tensor, + k_grad_out: torch.Tensor, + v_grad_out: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + qkv_split_arg_list: List[int], + qkv_format, + interleaved: bool, + cp_size: int, + cp_rank: int, +) -> torch.Tensor: + _require_triton() + _check_freqs(q_freqs, "q_freqs") + _check_freqs(k_freqs, "k_freqs") + if not q_freqs.is_contiguous(): + q_freqs = q_freqs.contiguous() + if not k_freqs.is_contiguous(): + k_freqs = k_freqs.contiguous() + q_grad_out = q_grad_out.contiguous() + k_grad_out = k_grad_out.contiguous() + v_grad_out = v_grad_out.contiguous() + + qkv_format = int(qkv_format) + is_sbhd = qkv_format == NVTE_SBHD + s = q_grad_out.size(0) if is_sbhd else q_grad_out.size(1) + b = q_grad_out.size(1) if is_sbhd else q_grad_out.size(0) + q_split, k_split, v_split = _check_qkv_splits(qkv_split_arg_list) + if q_grad_out.size(3) != k_split or k_grad_out.size(3) != k_split: + raise ValueError("Q and K gradient last dimensions must match the K split size") + if v_grad_out.size(3) != v_split: + raise ValueError("V gradient last dimension must match the V split size") + total_hd = (q_grad_out.size(2) + k_grad_out.size(2) + v_grad_out.size(2)) * q_grad_out.size(3) + total_d = q_split + k_split + v_split + if total_hd % total_d != 0: + raise ValueError("Q/K/V gradient shapes are inconsistent with qkv split sizes") + qkv_grad_size = list(q_grad_out.size()) + qkv_grad_size[2] = total_hd // total_d + qkv_grad_size[3] = total_d + h = qkv_grad_size[2] + d = v_split + d2 = q_freqs.size(3) + if d < d2: + raise ValueError("qkv value split must be greater than or equal to q_freqs last dim") + if q_freqs.size(3) != k_freqs.size(3): + raise ValueError("q_freqs and k_freqs must have the same rotary dimension") + + qkv_grad_input = torch.empty(qkv_grad_size, dtype=q_grad_out.dtype, device=q_grad_out.device) + block_h = _choose_qkv_block_h(h) + block_d = _choose_block_d(d) + d_blocks = triton.cdiv(d, block_d) + grid = (s, b, triton.cdiv(h, block_h) * d_blocks) + _fused_qkv_rope_kernel[grid]( + q_grad_out, + q_freqs, + k_freqs, + q_grad_out, + q_grad_out, + k_grad_out, + v_grad_out, + qkv_grad_input, + s, + b, + h, + d, + d2, + q_split, + k_split, + v_split, + qkv_format, + interleaved, + True, + False, + cp_size, + cp_rank, + BLOCK_H=block_h, + BLOCK_D=block_d, + N_D_BLOCKS=d_blocks, + num_warps=_num_warps(block_h), + ) + return qkv_grad_input diff --git a/transformer_engine/plugin/core/backends/flagos/register_ops.py b/transformer_engine/plugin/core/backends/flagos/register_ops.py index 01ae9610c7..f373e4cdfe 100644 --- a/transformer_engine/plugin/core/backends/flagos/register_ops.py +++ b/transformer_engine/plugin/core/backends/flagos/register_ops.py @@ -180,6 +180,39 @@ def register_builtins(registry) -> None: vendor=None, priority=150, ), + # RoPE (Rotary Position Embedding) + OpImpl( + op_name="fused_rope_forward", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available(backend.fused_rope_forward, is_avail), + vendor=None, + priority=150, + ), + OpImpl( + op_name="fused_rope_backward", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available(backend.fused_rope_backward, is_avail), + vendor=None, + priority=150, + ), + OpImpl( + op_name="fused_qkv_rope_forward", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available(backend.fused_qkv_rope_forward, is_avail), + vendor=None, + priority=150, + ), + OpImpl( + op_name="fused_qkv_rope_backward", + impl_id="default.flagos", + kind=BackendImplKind.DEFAULT, + fn=_bind_is_available(backend.fused_qkv_rope_backward, is_avail), + vendor=None, + priority=150, + ), OpImpl( op_name="get_cudnn_version", impl_id="default.flagos", diff --git a/transformer_engine/plugin/tests/run_all_tests.py b/transformer_engine/plugin/tests/run_all_tests.py index 1a0e02c615..ecd7d5be0d 100644 --- a/transformer_engine/plugin/tests/run_all_tests.py +++ b/transformer_engine/plugin/tests/run_all_tests.py @@ -11,6 +11,7 @@ from test_optimizer import OptimizerTests from test_flash_attention import FlashAttentionTests from test_te_general_grouped import grouped_gemmTests +from test_fused_rope import FusedRoPETests from test_policy import run_all_tests @@ -30,6 +31,7 @@ def main(): OptimizerTests(device=device), FlashAttentionTests(device=device), grouped_gemmTests(device=device), + FusedRoPETests(device=device), ] results = [] diff --git a/transformer_engine/plugin/tests/test_fused_rope.py b/transformer_engine/plugin/tests/test_fused_rope.py new file mode 100644 index 0000000000..d93aed642e --- /dev/null +++ b/transformer_engine/plugin/tests/test_fused_rope.py @@ -0,0 +1,766 @@ +# Copyright (c) 2025, BAAI. All rights reserved. +# +# See LICENSE for license information. + +from __future__ import annotations + +from typing import Optional + +import torch + +from transformer_engine.plugin.core.ops import NVTE_QKV_Format +from transformer_engine.plugin.test_utils import TestCase, get_available_backends, get_backend + + +def _triton_available() -> bool: + try: + import triton # noqa: F401 + except ModuleNotFoundError: + return False + return True + + +def _make_freqs(seq_len: int, d2: int, device: str) -> torch.Tensor: + values = torch.linspace(-0.7, 0.9, steps=seq_len * d2, dtype=torch.float32, device=device) + return values.reshape(seq_len, 1, 1, d2).contiguous() + + +def _freq_position( + s_id: int, + b_id: int, + cur_seqlens: int, + start_positions: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, +) -> int: + pos = s_id + if start_positions is not None: + pos += int(start_positions[b_id].item()) + + if cp_size > 1: + half = cur_seqlens // 2 + if s_id < half: + pos += cp_rank * half + else: + pos += cur_seqlens * cp_size - (cp_rank + 1) * half - half + return pos + + +def _apply_rope_slice( + src: torch.Tensor, + freq: torch.Tensor, + interleaved: bool, + is_backward: bool, +) -> torch.Tensor: + d2 = freq.numel() + out = src.clone() + src_rot = src[..., :d2].float() + + idx = torch.arange(d2, device=src.device) + if interleaved: + even = (idx % 2) == 0 + rot_idx = torch.where(even, idx + 1, idx - 1) + if is_backward: + sin_idx = rot_idx + sin_sign = torch.where(even, 1.0, -1.0) + rot_sign = torch.ones_like(freq) + else: + sin_idx = idx + sin_sign = torch.ones_like(freq) + rot_sign = torch.where(even, -1.0, 1.0) + else: + half = d2 // 2 + first_half = (idx + half) < d2 + rot_idx = torch.where(first_half, idx + half, idx + half - d2) + if is_backward: + sin_idx = rot_idx + sin_sign = torch.where(first_half, 1.0, -1.0) + rot_sign = torch.ones_like(freq) + else: + sin_idx = idx + sin_sign = torch.ones_like(freq) + rot_sign = torch.where(first_half, -1.0, 1.0) + + rotary = ( + src_rot * torch.cos(freq) + + src_rot[..., rot_idx] * rot_sign * torch.sin(freq[sin_idx]) * sin_sign + ) + out[..., :d2] = rotary.to(src.dtype) + return out + + +def _reference_rope( + tensor: torch.Tensor, + freqs: torch.Tensor, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + start_positions: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + is_backward: bool, +) -> torch.Tensor: + freq_flat = freqs[:, 0, 0, :] + out = torch.empty(tensor.size(), dtype=tensor.dtype, device=tensor.device) + + if qkv_format == NVTE_QKV_Format.NVTE_THD: + cu = (cu_seqlens.cpu() // cp_size).tolist() + for b_id in range(len(cu) - 1): + start, end = cu[b_id], cu[b_id + 1] + cur_seqlens = end - start + for s_id in range(cur_seqlens): + t_id = start + s_id + pos = _freq_position(s_id, b_id, cur_seqlens, start_positions, cp_size, cp_rank) + out[t_id] = _apply_rope_slice( + tensor[t_id], freq_flat[pos], interleaved, is_backward + ) + return out + + if qkv_format == NVTE_QKV_Format.NVTE_SBHD: + s, b = tensor.size(0), tensor.size(1) + for s_id in range(s): + for b_id in range(b): + pos = _freq_position(s_id, b_id, s, start_positions, cp_size, cp_rank) + out[s_id, b_id] = _apply_rope_slice( + tensor[s_id, b_id], freq_flat[pos], interleaved, is_backward + ) + return out + + s, b = tensor.size(1), tensor.size(0) + for b_id in range(b): + for s_id in range(s): + pos = _freq_position(s_id, b_id, s, start_positions, cp_size, cp_rank) + out[b_id, s_id] = _apply_rope_slice( + tensor[b_id, s_id], freq_flat[pos], interleaved, is_backward + ) + return out + + +def _reference_qkv_forward( + qkv: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + start_positions: Optional[torch.Tensor], + qkv_split_arg_list, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cp_size: int, + cp_rank: int, +): + q_split, k_split, v_split = qkv_split_arg_list + d = v_split + is_sbhd = qkv_format == NVTE_QKV_Format.NVTE_SBHD + s = qkv.size(0) if is_sbhd else qkv.size(1) + b = qkv.size(1) if is_sbhd else qkv.size(0) + h = qkv.size(2) + + q_out_size = list(qkv.size()) + q_out_size[2] = q_out_size[2] * q_split // k_split + q_out_size[3] = k_split + k_out_size = list(qkv.size()) + k_out_size[3] = k_split + v_out_size = list(qkv.size()) + v_out_size[3] = v_split + + q_out = torch.empty(q_out_size, dtype=qkv.dtype, device=qkv.device) + k_out = torch.empty(k_out_size, dtype=qkv.dtype, device=qkv.device) + v_out = torch.empty(v_out_size, dtype=qkv.dtype, device=qkv.device) + q_freq_flat = q_freqs[:, 0, 0, :] + k_freq_flat = k_freqs[:, 0, 0, :] + + for s_id in range(s): + for b_id in range(b): + pos = _freq_position(s_id, b_id, s, start_positions, cp_size, cp_rank) + src = qkv[s_id, b_id] if is_sbhd else qkv[b_id, s_id] + q_flat = (q_out[s_id, b_id] if is_sbhd else q_out[b_id, s_id]).reshape(-1) + k_flat = (k_out[s_id, b_id] if is_sbhd else k_out[b_id, s_id]).reshape(-1) + v_flat = (v_out[s_id, b_id] if is_sbhd else v_out[b_id, s_id]).reshape(-1) + + for h_id in range(h): + for row_offset in range(0, q_split, d): + q_slice = src[h_id, row_offset : row_offset + d] + q_flat[h_id * q_split + row_offset : h_id * q_split + row_offset + d] = ( + _apply_rope_slice(q_slice, q_freq_flat[pos], interleaved, False) + ) + k_start = q_split + for row_offset in range(0, k_split, d): + k_slice = src[h_id, k_start + row_offset : k_start + row_offset + d] + k_flat[h_id * k_split + row_offset : h_id * k_split + row_offset + d] = ( + _apply_rope_slice(k_slice, k_freq_flat[pos], interleaved, False) + ) + v_start = q_split + k_split + v_flat[h_id * v_split : (h_id + 1) * v_split] = src[ + h_id, v_start : v_start + v_split + ] + + return q_out, k_out, v_out + + +def _reference_qkv_backward( + q_grad: torch.Tensor, + k_grad: torch.Tensor, + v_grad: torch.Tensor, + q_freqs: torch.Tensor, + k_freqs: torch.Tensor, + qkv_split_arg_list, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cp_size: int, + cp_rank: int, +) -> torch.Tensor: + q_split, k_split, v_split = qkv_split_arg_list + d = v_split + total_d = q_split + k_split + v_split + total_hd = (q_grad.size(2) + k_grad.size(2) + v_grad.size(2)) * q_grad.size(3) + qkv_grad_size = list(q_grad.size()) + qkv_grad_size[2] = total_hd // total_d + qkv_grad_size[3] = total_d + out = torch.empty(qkv_grad_size, dtype=q_grad.dtype, device=q_grad.device) + + is_sbhd = qkv_format == NVTE_QKV_Format.NVTE_SBHD + s = q_grad.size(0) if is_sbhd else q_grad.size(1) + b = q_grad.size(1) if is_sbhd else q_grad.size(0) + h = out.size(2) + q_freq_flat = q_freqs[:, 0, 0, :] + k_freq_flat = k_freqs[:, 0, 0, :] + + for s_id in range(s): + for b_id in range(b): + pos = _freq_position(s_id, b_id, s, None, cp_size, cp_rank) + q_flat = (q_grad[s_id, b_id] if is_sbhd else q_grad[b_id, s_id]).reshape(-1) + k_flat = (k_grad[s_id, b_id] if is_sbhd else k_grad[b_id, s_id]).reshape(-1) + v_flat = (v_grad[s_id, b_id] if is_sbhd else v_grad[b_id, s_id]).reshape(-1) + dst = out[s_id, b_id] if is_sbhd else out[b_id, s_id] + + for h_id in range(h): + for row_offset in range(0, q_split, d): + q_slice = q_flat[h_id * q_split + row_offset : h_id * q_split + row_offset + d] + dst[h_id, row_offset : row_offset + d] = _apply_rope_slice( + q_slice, q_freq_flat[pos], interleaved, True + ) + k_start = q_split + for row_offset in range(0, k_split, d): + k_slice = k_flat[h_id * k_split + row_offset : h_id * k_split + row_offset + d] + dst[h_id, k_start + row_offset : k_start + row_offset + d] = _apply_rope_slice( + k_slice, k_freq_flat[pos], interleaved, True + ) + v_start = q_split + k_split + dst[h_id, v_start : v_start + v_split] = v_flat[ + h_id * v_split : (h_id + 1) * v_split + ] + + return out + + +class _TorchRoPEBackend: + @staticmethod + def fused_rope_forward( + input, + freqs, + start_positions, + qkv_format, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ): + return _reference_rope( + input, + freqs, + qkv_format, + interleaved, + cu_seqlens, + start_positions, + cp_size, + cp_rank, + False, + ) + + @staticmethod + def fused_rope_backward( + output_grads, + freqs, + start_positions, + qkv_format, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ): + return _reference_rope( + output_grads, + freqs, + qkv_format, + interleaved, + cu_seqlens, + start_positions, + cp_size, + cp_rank, + True, + ) + + @staticmethod + def fused_qkv_rope_forward( + qkv_input, + q_freqs, + k_freqs, + start_positions, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ): + return _reference_qkv_forward( + qkv_input, + q_freqs, + k_freqs, + start_positions, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + + @staticmethod + def fused_qkv_rope_backward( + q_grad_out, + k_grad_out, + v_grad_out, + q_freqs, + k_freqs, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ): + return _reference_qkv_backward( + q_grad_out, + k_grad_out, + v_grad_out, + q_freqs, + k_freqs, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + + +class FusedRoPETests(TestCase): + def __init__(self, device="cpu"): + super().__init__( + "Fused RoPE", + "Test fused RoPE and fused QKV RoPE across CUDA, FlagOS, and torch reference", + ) + self.backends = get_available_backends() + if "torch" not in self.backends: + self.backends.append("torch") + self.backends = [ + backend for backend in self.backends if backend in ("cuda", "flagos", "torch") + ] + self.device = device + + def _get_backend(self, backend_name): + if backend_name == "torch": + return _TorchRoPEBackend() + if self.device == "cpu": + raise NotImplementedError("fused RoPE requires a GPU device") + if backend_name == "flagos" and not _triton_available(): + raise NotImplementedError("Triton is not installed") + return get_backend(backend_name) + + def _iter_backends(self): + if not self.backends: + self.skipped += 1 + print(" ⊘ no tested backend is registered") + return + for backend_name in self.backends: + try: + yield backend_name, self._get_backend(backend_name) + except NotImplementedError as exc: + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + + def _compare_to_cuda(self, outputs, backend_name, labels, msg): + if "cuda" not in outputs or backend_name not in outputs: + return + + try: + for actual, expected, label in zip(outputs[backend_name], outputs["cuda"], labels): + self.assert_close( + actual.float(), + expected.float(), + rtol=1e-4, + atol=1e-4, + msg=f"{msg} {label} mismatch between {backend_name} and cuda", + ) + print(f" ✓ {backend_name} matches cuda") + except AssertionError as exc: + print(f" ✗ {backend_name} vs cuda: {exc}") + + def test_rope_sbhd_bshd_forward_backward(self): + print("\n Testing fused_rope_forward/backward for SBHD and BSHD") + cases = [ + (NVTE_QKV_Format.NVTE_SBHD, (5, 2, 3, 10), False, 1, 0, True), + (NVTE_QKV_Format.NVTE_BSHD, (2, 4, 2, 10), True, 2, 1, False), + (NVTE_QKV_Format.NVTE_SBHD, (4, 2, 2, 10), False, 2, 1, True), + ] + + for qkv_format, shape, interleaved, cp_size, cp_rank, use_start in cases: + print( + f"\n Testing fused_rope_forward/backward with {qkv_format.name}, " + f"interleaved={interleaved}, cp_size={cp_size}, " + f"start_positions={use_start}" + ) + d2 = 6 + freq_len = shape[0] if qkv_format == NVTE_QKV_Format.NVTE_SBHD else shape[1] + freq_len = max(freq_len * cp_size + 3, 12) + freqs = _make_freqs(freq_len, d2, self.device) + start_positions = None + if use_start: + batch = shape[1] if qkv_format == NVTE_QKV_Format.NVTE_SBHD else shape[0] + start_positions = torch.arange(batch, dtype=torch.int32, device=self.device) + 1 + + base = torch.randn(*shape[:-1], shape[-1] * 2, device=self.device) + tensor = base[..., ::2] + grad = torch.randn_like(tensor) + ref_fwd = _reference_rope( + tensor, + freqs, + qkv_format, + interleaved, + None, + start_positions, + cp_size, + cp_rank, + False, + ) + ref_bwd = _reference_rope( + grad, freqs, qkv_format, interleaved, None, start_positions, cp_size, cp_rank, True + ) + + outputs = {} + for backend_name, backend in self._iter_backends(): + try: + out = backend.fused_rope_forward( + tensor, + freqs, + start_positions, + qkv_format, + interleaved, + None, + cp_size, + cp_rank, + ) + dx = backend.fused_rope_backward( + grad, + freqs, + start_positions, + qkv_format, + interleaved, + None, + cp_size, + cp_rank, + ) + self.assert_close( + out.float(), + ref_fwd.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_rope_forward mismatch for {backend_name}", + ) + self.assert_close( + dx.float(), + ref_bwd.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_rope_backward mismatch for {backend_name}", + ) + outputs[backend_name] = (out, dx) + print(f" ✓ {backend_name}") + except NotImplementedError as exc: + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + except RuntimeError as exc: + if "is not available" in str(exc): + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + else: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + except Exception as exc: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + + self._compare_to_cuda( + outputs, + "flagos", + ("forward", "backward"), + f"{qkv_format.name} fused_rope", + ) + + def test_rope_thd_forward_backward(self): + print("\n Testing fused_rope_forward/backward for THD") + cases = [ + (torch.tensor([0, 3, 8], dtype=torch.int32), True, 1, 0, True), + (torch.tensor([0, 8, 20], dtype=torch.int32), False, 2, 0, False), + ] + + for cu_cpu, interleaved, cp_size, cp_rank, use_start in cases: + print( + "\n Testing fused_rope_forward/backward with NVTE_THD, " + f"interleaved={interleaved}, cp_size={cp_size}, " + f"start_positions={use_start}" + ) + cu_seqlens = cu_cpu.to(self.device) + local_cu = cu_cpu // cp_size + total_t = int(local_cu[-1].item()) + h, d, d2 = 3, 10, 6 + freq_len = max(int(cu_cpu[1:].sub(cu_cpu[:-1]).max().item()), 12) + freqs = _make_freqs(freq_len, d2, self.device) + start_positions = None + if use_start: + start_positions = torch.tensor([1, 0], dtype=torch.int32, device=self.device) + + tensor = torch.randn(total_t, h, d, device=self.device) + grad = torch.randn_like(tensor) + ref_fwd = _reference_rope( + tensor, + freqs, + NVTE_QKV_Format.NVTE_THD, + interleaved, + cu_seqlens, + start_positions, + cp_size, + cp_rank, + False, + ) + ref_bwd = _reference_rope( + grad, + freqs, + NVTE_QKV_Format.NVTE_THD, + interleaved, + cu_seqlens, + start_positions, + cp_size, + cp_rank, + True, + ) + + outputs = {} + for backend_name, backend in self._iter_backends(): + try: + out = backend.fused_rope_forward( + tensor, + freqs, + start_positions, + NVTE_QKV_Format.NVTE_THD, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ) + dx = backend.fused_rope_backward( + grad, + freqs, + start_positions, + NVTE_QKV_Format.NVTE_THD, + interleaved, + cu_seqlens, + cp_size, + cp_rank, + ) + self.assert_close( + out.float(), + ref_fwd.float(), + rtol=1e-4, + atol=1e-4, + msg=f"THD fused_rope_forward mismatch for {backend_name}", + ) + self.assert_close( + dx.float(), + ref_bwd.float(), + rtol=1e-4, + atol=1e-4, + msg=f"THD fused_rope_backward mismatch for {backend_name}", + ) + outputs[backend_name] = (out, dx) + print(f" ✓ {backend_name}") + except NotImplementedError as exc: + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + except RuntimeError as exc: + if "is not available" in str(exc): + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + else: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + except Exception as exc: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + + self._compare_to_cuda( + outputs, + "flagos", + ("forward", "backward"), + "THD fused_rope", + ) + + def test_qkv_rope_forward_backward(self): + print("\n Testing fused_qkv_rope_forward/backward") + cases = [ + (NVTE_QKV_Format.NVTE_SBHD, (4, 2, 2, 32), False, 1, 0, True), + (NVTE_QKV_Format.NVTE_BSHD, (2, 4, 2, 32), True, 2, 1, False), + ] + qkv_split_arg_list = [16, 8, 8] + + for qkv_format, shape, interleaved, cp_size, cp_rank, use_start in cases: + print( + f"\n Testing fused_qkv_rope_forward/backward with {qkv_format.name}, " + f"interleaved={interleaved}, cp_size={cp_size}, " + f"start_positions={use_start}" + ) + d2 = 6 + seq_len = shape[0] if qkv_format == NVTE_QKV_Format.NVTE_SBHD else shape[1] + freq_len = max(seq_len * cp_size + 3, 12) + q_freqs = _make_freqs(freq_len, d2, self.device) + k_freqs = _make_freqs(freq_len, d2, self.device) + 0.17 + start_positions = None + if use_start: + batch = shape[1] if qkv_format == NVTE_QKV_Format.NVTE_SBHD else shape[0] + start_positions = torch.arange(batch, dtype=torch.int32, device=self.device) + + qkv = torch.randn(*shape, device=self.device).contiguous() + ref_q, ref_k, ref_v = _reference_qkv_forward( + qkv, + q_freqs, + k_freqs, + start_positions, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + + q_grad = torch.randn_like(ref_q) + k_grad = torch.randn_like(ref_k) + v_grad = torch.randn_like(ref_v) + ref_bwd = _reference_qkv_backward( + q_grad, + k_grad, + v_grad, + q_freqs, + k_freqs, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + + outputs = {} + for backend_name, backend in self._iter_backends(): + try: + q_out, k_out, v_out = backend.fused_qkv_rope_forward( + qkv, + q_freqs, + k_freqs, + start_positions, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + dqkv = backend.fused_qkv_rope_backward( + q_grad, + k_grad, + v_grad, + q_freqs, + k_freqs, + qkv_split_arg_list, + qkv_format, + interleaved, + cp_size, + cp_rank, + ) + self.assert_close( + q_out.float(), + ref_q.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_qkv_rope_forward Q mismatch for {backend_name}", + ) + self.assert_close( + k_out.float(), + ref_k.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_qkv_rope_forward K mismatch for {backend_name}", + ) + self.assert_close( + v_out.float(), + ref_v.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_qkv_rope_forward V mismatch for {backend_name}", + ) + self.assert_close( + dqkv.float(), + ref_bwd.float(), + rtol=1e-4, + atol=1e-4, + msg=f"fused_qkv_rope_backward mismatch for {backend_name}", + ) + outputs[backend_name] = (q_out, k_out, v_out, dqkv) + print(f" ✓ {backend_name}") + except NotImplementedError as exc: + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + except RuntimeError as exc: + if "is not available" in str(exc): + self.skipped += 1 + print(f" ⊘ {backend_name} ({exc})") + else: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + except Exception as exc: + self.failed += 1 + print(f" ✗ {backend_name}: {exc}") + + self._compare_to_cuda( + outputs, + "flagos", + ("Q forward", "K forward", "V forward", "backward"), + f"{qkv_format.name} fused_qkv_rope", + ) + + def run_all_tests(self): + print("\n" + "=" * 60) + print("Testing Fused RoPE") + print("=" * 60) + print(f"Available backends: {', '.join(self.backends) or 'none'}") + + self.test_rope_sbhd_bshd_forward_backward() + self.test_rope_thd_forward_backward() + self.test_qkv_rope_forward_backward() + + return self.report() + + +def main(): + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {device}") + test_suite = FusedRoPETests(device=device) + success = test_suite.run_all_tests() + return 0 if success else 1 + + +if __name__ == "__main__": + exit(main())