diff --git a/transformer_engine/plugin/core/backends/flagos/flagos.py b/transformer_engine/plugin/core/backends/flagos/flagos.py index ecdc73b33a..03f7c2ed7e 100644 --- a/transformer_engine/plugin/core/backends/flagos/flagos.py +++ b/transformer_engine/plugin/core/backends/flagos/flagos.py @@ -7,7 +7,7 @@ import torch -from ...ops import TEFLBackendBase, FP8TensorMeta, NVTE_Fused_Attn_Backend +from ...ops import * from .impl import ( rmsnorm_fwd_fl, rmsnorm_bwd_fl, @@ -20,7 +20,6 @@ def _check_flagos_available() -> bool: return True - class FlagOSBackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -29,10 +28,6 @@ def check_available() -> bool: def is_available(self) -> bool: return _check_flagos_available() - def get_flash_attention_class(self): - from .attention.dot_product_attention.backends import FlashAttentionFL - return FlashAttentionFL - def get_attention_backend(self, attention_params=None): from packaging.version import Version as PkgVersion from ...logger_manager import get_logger @@ -65,17 +60,18 @@ def get_attention_backend(self, attention_params=None): available_backends, ) +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -84,53 +80,53 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: + ) -> List[Any]: return generic_gemm_fl( A, transA, B, transB, D, quantizer, output_dtype, - bias, bias_type, gelu, gelu_in, grad, - workspace, workspace_size, accumulate, use_split_accumulator, - comm_overlap=comm_overlap, comm_type=comm_type, - extra_output=extra_output, bulk_overlap=bulk_overlap, - alpha=alpha, beta=beta + bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, + accumulate, use_split_accumulator, comm_overlap, comm_type, + extra_output, bulk_overlap, alpha, beta ) + # Other granular functions def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: return rmsnorm_fwd_fl( input=input, weight=weight, eps=eps, ln_out=ln_out, quantizer=quantizer, odtype=otype, sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma, ) - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: return rmsnorm_bwd_fl( - dy=dy, x=x, rsigma=rsigma, gamma=gamma, - sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma, eps=eps, + dy=dz, x=x, rsigma=rsigma, gamma=gamma, + sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma ) + def get_fused_attn_backend(self, *args, **kwargs) -> int: + return NVTE_Fused_Attn_Backend.NVTE_No_Backend + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -139,73 +135,61 @@ def multi_tensor_scale( scale: float, ) -> None: return multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale) - def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: - result, _ = multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor) - return result - + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: + return multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor) def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): - if chunk_size is None: - return multi_tensor_adam_fl + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: return multi_tensor_adam_fl( - chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists, - lr=lr, beta1=beta1, beta2=beta2, eps=eps, - step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay, + chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, ) - def multi_tensor_adam_param_remainder( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): - if chunk_size is None: - return multi_tensor_adam_param_remainder_fl + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: return multi_tensor_adam_param_remainder_fl( - chunk_size=chunk_size, noop_flag=noop_flag, tensor_lists=tensor_lists, - lr=lr, beta1=beta1, beta2=beta2, eps=eps, - step=step, mode=mode, bias_correction=bias_correction, weight_decay=weight_decay, + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, ) + # Misc def get_cublasLt_version(self) -> int: return 110000 - def get_cudnn_version(self) -> int: return 90000 - def get_num_cublas_streams(self) -> int: return 0 - def get_fused_attn_backend(self, *args, **kwargs) -> int: - return NVTE_Fused_Attn_Backend.NVTE_No_Backend - - def create_fp8_tensor_meta(self) -> FP8TensorMeta: - return FP8TensorMeta() - +############## class func ################################# + def get_flash_attention_class(self): + from .attention.dot_product_attention.backends import FlashAttentionFL + return FlashAttentionFL diff --git a/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py b/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py index 93ba067e93..89107b04c2 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/fused_adam.py @@ -187,4 +187,4 @@ def multi_tensor_adam_param_remainder_fl( # Write back flag_gems.copy_(p, param_bf16) - flag_gems.copy_(p_remainder, remainder_int16) + flag_gems.copy_(p_remainder, remainder_int16) \ No newline at end of file diff --git a/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py b/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py index 4421487ff1..d7361fd7ed 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/multi_tensor.py @@ -23,4 +23,4 @@ def multi_tensor_l2_norm_fl(chunk_size, noop_flag, tensor_lists, per_tensor, *ar def multi_tensor_scale_fl(chunk_size, noop_flag, tensor_lists, scale): for src, dst in zip(tensor_lists[0], tensor_lists[1]): - flag_gems.copy_(dst, src * scale) + flag_gems.copy_(dst, src * scale) \ No newline at end of file diff --git a/transformer_engine/plugin/core/backends/flagos/impl/rmsnorm.py b/transformer_engine/plugin/core/backends/flagos/impl/rmsnorm.py index ffa382147f..12fda567ed 100644 --- a/transformer_engine/plugin/core/backends/flagos/impl/rmsnorm.py +++ b/transformer_engine/plugin/core/backends/flagos/impl/rmsnorm.py @@ -42,7 +42,7 @@ def rmsnorm_bwd_fl( gamma, sm_margin, zero_centered_gamma, - eps, + eps=1e-5, ): # When zero_centered_gamma is True, forward uses (1 + gamma) as weight # So backward needs to use (1 + gamma) for computing dx diff --git a/transformer_engine/plugin/core/backends/reference/impl/normalization.py b/transformer_engine/plugin/core/backends/reference/impl/normalization.py index 6ab7a7648c..48f89b44d8 100644 --- a/transformer_engine/plugin/core/backends/reference/impl/normalization.py +++ b/transformer_engine/plugin/core/backends/reference/impl/normalization.py @@ -5,12 +5,37 @@ from typing import Any, Optional, Tuple import torch import torch.nn.functional as F +from ....ops import DType __all__ = [ "layernorm_fwd_torch", "layernorm_bwd_torch", ] +# Mapping from DType enum to torch.dtype +_DTYPE_TO_TORCH_DTYPE = { + DType.kByte: torch.uint8, + DType.kInt16: torch.int16, + DType.kInt32: torch.int32, + DType.kInt64: torch.int64, + DType.kFloat32: torch.float32, + DType.kFloat16: torch.float16, + DType.kBFloat16: torch.bfloat16, + DType.kFloat8E4M3: torch.float8_e4m3fn, + DType.kFloat8E5M2: torch.float8_e5m2, +} + +def _to_torch_dtype(dtype): + """Convert DType enum to torch.dtype.""" + if dtype is None: + return None + if isinstance(dtype, torch.dtype): + return dtype + if isinstance(dtype, (int, DType)): + dtype_enum = DType(dtype) + if dtype_enum in _DTYPE_TO_TORCH_DTYPE: + return _DTYPE_TO_TORCH_DTYPE[dtype_enum] + raise ValueError(f"Unsupported dtype: {dtype}") def layernorm_fwd_torch( input: torch.Tensor, @@ -19,10 +44,11 @@ def layernorm_fwd_torch( eps: float, ln_out: Optional[torch.Tensor], quantizer: Any, - odtype: torch.dtype, + odtype: DType, sm_margin: int, zero_centered_gamma: bool, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + odtype = _to_torch_dtype(odtype) mean = input.mean(dim=-1, keepdim=True) var = input.var(dim=-1, keepdim=True, unbiased=False) rsigma = torch.rsqrt(var + eps) @@ -45,7 +71,6 @@ def layernorm_fwd_torch( return output, mean, rsigma - def layernorm_bwd_torch( dy: torch.Tensor, x: torch.Tensor, diff --git a/transformer_engine/plugin/core/backends/reference/impl/optimizer.py b/transformer_engine/plugin/core/backends/reference/impl/optimizer.py index 0ae0809dcc..f3140a5695 100644 --- a/transformer_engine/plugin/core/backends/reference/impl/optimizer.py +++ b/transformer_engine/plugin/core/backends/reference/impl/optimizer.py @@ -310,4 +310,4 @@ def multi_tensor_compute_scale_and_scale_inv_torch( # Update scale and scale_inv scale.copy_(computed_scale) - scale_inv.copy_(1.0 / computed_scale) + scale_inv.copy_(1.0 / computed_scale) \ No newline at end of file diff --git a/transformer_engine/plugin/core/backends/reference/impl/rmsnorm.py b/transformer_engine/plugin/core/backends/reference/impl/rmsnorm.py index 7ae420e7f3..0aebdae2fe 100644 --- a/transformer_engine/plugin/core/backends/reference/impl/rmsnorm.py +++ b/transformer_engine/plugin/core/backends/reference/impl/rmsnorm.py @@ -43,7 +43,6 @@ def rmsnorm_bwd_torch( gamma, sm_margin, zero_centered_gamma, - eps, ): inv_rms = rsigma.unsqueeze(-1) diff --git a/transformer_engine/plugin/core/backends/reference/reference.py b/transformer_engine/plugin/core/backends/reference/reference.py index 3f29cf89be..80c7b327f0 100644 --- a/transformer_engine/plugin/core/backends/reference/reference.py +++ b/transformer_engine/plugin/core/backends/reference/reference.py @@ -3,11 +3,9 @@ # See LICENSE for license information. import os -from typing import Any, Dict, List, Optional, Tuple, Union - +from typing import Any, List, Optional, Tuple import torch - -from ...ops import TEFLBackendBase, FP8TensorMeta, NVTE_Fused_Attn_Backend +from ...ops import * from .impl import ( general_gemm_torch, @@ -33,6 +31,7 @@ multi_tensor_sgd_torch, ) + class ReferenceBackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -41,11 +40,7 @@ def check_available() -> bool: def is_available(self) -> bool: return True - def get_flash_attention_class(self): - from .flash_attention import FlashAttentionTorch - return FlashAttentionTorch - - def get_attention_backend(self, attention_params=None): + def get_attention_backend(self, _attention_params=None): from packaging.version import Version as PkgVersion from ...logger_manager import get_logger logger = get_logger() @@ -79,15 +74,15 @@ def get_attention_backend(self, attention_params=None): def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -96,49 +91,20 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: + ) -> List[Any]: return general_gemm_torch( - A=A, - transA=transA, - B=B, - transB=transB, - D=D, - quantizer=quantizer, - output_dtype=output_dtype, - bias=bias, - bias_type=bias_type, - gelu=gelu, - gelu_in=gelu_in, - grad=grad, - workspace=workspace, - workspace_size=workspace_size, - accumulate=accumulate, - use_split_accumulator=use_split_accumulator, - comm_overlap=comm_overlap, - comm_type=comm_type, - extra_output=extra_output, - bulk_overlap=bulk_overlap, - alpha=alpha, - beta=beta, + A, transA, B, transB, D, quantizer, output_dtype, + bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, + accumulate, use_split_accumulator, comm_overlap, comm_type, + extra_output, bulk_overlap, alpha, beta ) - def te_general_grouped_gemm(self, *args, **kwargs) -> Any: - raise NotImplementedError("te_general_grouped_gemm - not implemented in reference backend") - - def quantize(self, tensor: torch.Tensor, quantizer: Any, output: Optional[torch.Tensor] = None, noop: Optional[torch.Tensor] = None) -> Any: - raise NotImplementedError("quantize - not implemented in reference backend") - - def dequantize(self, input: torch.Tensor, otype: torch.dtype) -> torch.Tensor: - raise NotImplementedError("dequantize - not implemented in reference backend") - - def bgrad_quantize(self, input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: - raise NotImplementedError("bgrad_quantize - not implemented in reference backend") - + # GELU and variants def gelu(self, input: torch.Tensor, quantizer: Any) -> Any: return gelu_torch(input, quantizer) @@ -151,6 +117,7 @@ def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any: def qgeglu(self, input: torch.Tensor, quantizer: Any) -> Any: return qgeglu_torch(input, quantizer) + # ReLU and variants def relu(self, input: torch.Tensor, quantizer: Any) -> Any: return relu_torch(input, quantizer) @@ -163,15 +130,23 @@ def srelu(self, input: torch.Tensor, quantizer: Any) -> Any: def sreglu(self, input: torch.Tensor, quantizer: Any) -> Any: return sreglu_torch(input, quantizer) + # SwiGLU and variants def silu(self, input: torch.Tensor, quantizer: Any) -> Any: return silu_torch(input, quantizer) def swiglu(self, input: torch.Tensor, quantizer: Any) -> Any: return swiglu_torch(input, quantizer) - def clamped_swiglu(self, input: torch.Tensor, quantizer: Any, limit: float = 7.0, alpha: float = 1.702) -> Any: + def clamped_swiglu( + self, + input: torch.Tensor, + quantizer: Any, + limit: float = 7.0, + alpha: float = 1.702, + ) -> Any: return clamped_swiglu_torch(input, quantizer, limit, alpha) + # Backward of GELU and variants def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return dgelu_torch(grad, fwd_input, quantizer) @@ -184,6 +159,7 @@ def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> def dqgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return dqgeglu_torch(grad, fwd_input, quantizer) + # Backward of ReLU and variants def drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return drelu_torch(grad, fwd_input, quantizer) @@ -196,42 +172,77 @@ def dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> def dsreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return dsreglu_torch(grad, fwd_input, quantizer) + # Backward of SiLU and variants def dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return dsilu_torch(grad, fwd_input, quantizer) def dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: return dswiglu_torch(grad, fwd_input, quantizer) - def clamped_dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, limit: float = 7.0, alpha: float = 1.702) -> Any: + def clamped_dswiglu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + limit: float = 7.0, + alpha: float = 1.702, + ) -> Any: return clamped_dswiglu_torch(grad, fwd_input, quantizer, limit, alpha) - def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + # DBias + DAct fusions + def dbias_dgelu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + ) -> List[Any]: return dbias_dgelu_torch(grad, fwd_input, quantizer) - def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsilu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + ) -> List[Any]: return dbias_dsilu_torch(grad, fwd_input, quantizer) - def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_drelu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + ) -> Tuple[torch.Tensor, Any]: return dbias_drelu_torch(grad, fwd_input, quantizer) - def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dqgelu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + ) -> List[Any]: return dbias_dqgelu_torch(grad, fwd_input, quantizer) - def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsrelu( + self, + grad: torch.Tensor, + fwd_input: torch.Tensor, + quantizer: Any, + ) -> List[Any]: return dbias_dsrelu_torch(grad, fwd_input, quantizer) + # LayerNorm def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> List[Any]: return layernorm_fwd_torch( input=input, weight=weight, @@ -246,16 +257,16 @@ def layernorm_fwd( def layernorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: return layernorm_bwd_torch( - dy=dy, + dy=dz, x=x, mu=mu, rsigma=rsigma, @@ -264,17 +275,18 @@ def layernorm_bwd( zero_centered_gamma=zero_centered_gamma, ) + # RMSNorm def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: return rmsnorm_fwd_torch( input=input, weight=weight, @@ -288,153 +300,126 @@ def rmsnorm_fwd( def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: return rmsnorm_bwd_torch( - dy=dy, + dy=dz, x=x, rsigma=rsigma, gamma=gamma, sm_margin=sm_margin, zero_centered_gamma=zero_centered_gamma, - eps=eps, ) - def rmsnorm_bwd_add(self, *args, **kwargs) -> Any: - raise NotImplementedError("rmsnorm_bwd_add - not implemented in reference backend") - - def multi_tensor_quantize(self, tensor_list: List[torch.Tensor], quantizer_list: List[Any]) -> List[Any]: - raise NotImplementedError("multi_tensor_quantize - not implemented in reference backend") - - def split_quantize(self, tensor: torch.Tensor, split_sections: List[int], quantizer_list: List[Any]) -> List[Any]: - raise NotImplementedError("split_quantize - not implemented in reference backend") - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("moe_permute_fwd - not implemented in reference backend") - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("moe_permute_bwd - not implemented in reference backend") - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("moe_unpermute_fwd - not implemented in reference backend") - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("moe_unpermute_bwd - not implemented in reference backend") - - def scaled_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: return scaled_softmax_forward_torch(input, scale) - def scaled_softmax_backward(self, output_grad: torch.Tensor, softmax_output: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_softmax_backward_torch(output_grad, softmax_output, scale) - - def scaled_masked_softmax_forward(self, input: torch.Tensor, mask: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_masked_softmax_forward_torch(input, mask, scale) + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_softmax_backward_torch(output_grad_, softmax_results_, scale_factor) - def scaled_masked_softmax_backward(self, output_grad: torch.Tensor, softmax_output: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_masked_softmax_backward_torch(output_grad, softmax_output, scale) + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_masked_softmax_forward_torch(input, mask, scale_factor) - def scaled_upper_triang_masked_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_upper_triang_masked_softmax_forward_torch(input, scale) + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_masked_softmax_backward_torch(output_grad_, softmax_results_, scale_factor) - def scaled_upper_triang_masked_softmax_backward(self, output_grad: torch.Tensor, softmax_output: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_upper_triang_masked_softmax_backward_torch(output_grad, softmax_output, scale) + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_upper_triang_masked_softmax_forward_torch(input, scale_factor) - def scaled_aligned_causal_masked_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_aligned_causal_masked_softmax_forward_torch(input, scale) + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_upper_triang_masked_softmax_backward_torch(output_grads_, softmax_results_, scale_factor) - def scaled_aligned_causal_masked_softmax_backward(self, output_grad: torch.Tensor, softmax_output: torch.Tensor, scale: float) -> torch.Tensor: - return scaled_aligned_causal_masked_softmax_backward_torch(output_grad, softmax_output, scale) + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_aligned_causal_masked_softmax_forward_torch(input, scale_factor) - def get_fused_attn_backend(self, *args, **kwargs) -> int: + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + return scaled_aligned_causal_masked_softmax_backward_torch(output_grad_, softmax_results_, scale_factor) + + # Fused attention backend + def get_fused_attn_backend( + self, + _is_training: bool, + _q_dtype: DType, + _kv_dtype: DType, + _qkv_layout: NVTE_QKV_Layout, + _bias_type: NVTE_Bias_Type, + _attn_mask_type: NVTE_Mask_Type, + _softmax_type: NVTE_Softmax_Type, + _p_dropout: float, + _num_attn_heads: int, + _num_gqa_groups: int, + _max_seqlen_q: int, + _max_seqlen_kv: int, + _head_dim_qk: int, + _head_dim_v: int, + _window_size_left: int, + _window_size_right: int, + _return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: return NVTE_Fused_Attn_Backend.NVTE_No_Backend - def fused_attn_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_attn_fwd - not implemented in reference backend") - - def fused_attn_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_attn_bwd - not implemented in reference backend") - - def fa_prepare_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fa_prepare_fwd - not implemented in reference backend") - - def fa_prepare_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fa_prepare_bwd - not implemented in reference backend") - - def copy_to_kv_cache(self, *args, **kwargs) -> Any: - raise NotImplementedError("copy_to_kv_cache - not implemented in reference backend") - - def convert_thd_to_bshd(self, *args, **kwargs) -> Any: - raise NotImplementedError("convert_thd_to_bshd - not implemented in reference backend") - - def convert_bshd_to_thd(self, *args, **kwargs) -> Any: - raise NotImplementedError("convert_bshd_to_thd - not implemented in reference backend") - - def fused_rope_forward(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_rope_forward - not implemented in reference backend") - - def fused_rope_backward(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_rope_backward - not implemented in reference backend") - - def fused_qkv_rope_forward(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_qkv_rope_forward - not implemented in reference backend") - - def fused_qkv_rope_backward(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_qkv_rope_backward - not implemented in reference backend") - - def fused_topk_with_score_function_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_topk_with_score_function_fwd - not implemented in reference backend") - - def fused_topk_with_score_function_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_topk_with_score_function_bwd - not implemented in reference backend") - - def fused_score_for_moe_aux_loss_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_score_for_moe_aux_loss_fwd - not implemented in reference backend") - - def fused_score_for_moe_aux_loss_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_score_for_moe_aux_loss_bwd - not implemented in reference backend") - - def fused_moe_aux_loss_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_moe_aux_loss_fwd - not implemented in reference backend") - - def fused_moe_aux_loss_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_moe_aux_loss_bwd - not implemented in reference backend") - - def dropout_fwd(self, input: torch.Tensor, dropout_probability: float, out: Optional[torch.Tensor] = None) -> Tuple[torch.Tensor, torch.Tensor]: + # Dropout + def dropout_fwd( + self, + input: torch.Tensor, + dropout_probability: float, + out: Optional[torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor]: return dropout_fwd_torch(input, dropout_probability, out) - def dropout_bwd(self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, grad_input: Optional[torch.Tensor] = None) -> torch.Tensor: + def dropout_bwd( + self, + grad_output: torch.Tensor, + mask: torch.Tensor, + dropout_probability: float, + grad_input: Optional[torch.Tensor], + ) -> torch.Tensor: return dropout_bwd_torch(grad_output, mask, dropout_probability, grad_input) - def fp8_transpose(self, input: torch.Tensor, dtype: Any, *, out: torch.Tensor) -> None: - raise NotImplementedError("fp8_transpose - not implemented in reference backend") - - def swap_first_dims(self, tensor: torch.Tensor, *, out: torch.Tensor) -> None: - raise NotImplementedError("swap_first_dims - not implemented in reference backend") - - def compute_amax(self, input: torch.Tensor, amax: torch.Tensor) -> None: - raise NotImplementedError("compute_amax - not implemented in reference backend") - - def fused_amax_and_scale_update_after_reduction(self, *args, **kwargs) -> None: - raise NotImplementedError("fused_amax_and_scale_update_after_reduction - not implemented in reference backend") - - def fp8_block_scaling_compute_partial_amax(self, *args, **kwargs) -> None: - raise NotImplementedError("fp8_block_scaling_compute_partial_amax - not implemented in reference backend") - - def fp8_block_scaling_partial_cast(self, *args, **kwargs) -> None: - raise NotImplementedError("fp8_block_scaling_partial_cast - not implemented in reference backend") - - def fused_multi_row_padding(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_multi_row_padding - not implemented in reference backend") - - def fused_multi_row_unpadding(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_multi_row_unpadding - not implemented in reference backend") - + # Misc def get_cublasLt_version(self) -> int: return 0 @@ -444,100 +429,101 @@ def get_cudnn_version(self) -> int: def get_num_cublas_streams(self) -> int: return 0 - def thd_read_half_tensor(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_read_half_tensor - not implemented in reference backend") - - def thd_second_half_lse_correction(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_second_half_lse_correction - not implemented in reference backend") - - def thd_read_second_half_lse(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_read_second_half_lse - not implemented in reference backend") - - def thd_out_correction(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_out_correction - not implemented in reference backend") - - def thd_grad_correction(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_grad_correction - not implemented in reference backend") - - def thd_get_partitioned_indices(self, *args, **kwargs) -> Any: - raise NotImplementedError("thd_get_partitioned_indices - not implemented in reference backend") - - def init_nvshmem_backend(self, *args, **kwargs) -> None: - raise NotImplementedError("init_nvshmem_backend - not implemented in reference backend") - - def create_nvshmem_tensor(self, *args, **kwargs) -> torch.Tensor: - raise NotImplementedError("create_nvshmem_tensor - not implemented in reference backend") - - def nvshmem_send_on_current_stream(self, *args, **kwargs) -> None: - raise NotImplementedError("nvshmem_send_on_current_stream - not implemented in reference backend") - - def nvshmem_wait_on_current_stream(self, *args, **kwargs) -> None: - raise NotImplementedError("nvshmem_wait_on_current_stream - not implemented in reference backend") - - def nvshmem_finalize(self) -> None: - raise NotImplementedError("nvshmem_finalize - not implemented in reference backend") - - def multi_tensor_scale(self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], scale: float) -> None: + # Multi-tensor functions + def multi_tensor_scale( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + scale: float, + ) -> None: return multi_tensor_scale_torch(chunk_size, noop_flag, tensor_lists, scale) - def multi_tensor_l2norm(self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], per_tensor: bool = False) -> Union[torch.Tensor, List[torch.Tensor]]: + def multi_tensor_l2norm( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: return multi_tensor_l2norm_torch(chunk_size, noop_flag, tensor_lists, per_tensor) - def multi_tensor_unscale_l2norm(self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], scale: torch.Tensor, per_tensor: bool = False) -> Union[torch.Tensor, List[torch.Tensor]]: - """Compute L2 norm after unscaling. - - Note: scale parameter is actually inv_scale (1/loss_scale). - Unscaling means multiplying by inv_scale (= dividing by loss_scale). - """ + def multi_tensor_unscale_l2norm( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: if noop_flag.item() != 0: - if per_tensor: - return [torch.tensor(0.0, device=t.device) for t in tensor_lists[0]] - else: - return torch.tensor(0.0, device=tensor_lists[0][0].device) + device = tensor_lists[0][0].device if tensor_lists and tensor_lists[0] else 'cpu' + return torch.tensor(0.0, device=device), torch.tensor(0.0, device=device) - # Multiply by inv_scale (scale parameter is actually inverse scale) + # Multiply by inv_scale unscaled_tensors = [] for tensor in tensor_lists[0]: - unscaled_tensors.append(tensor * scale.item()) + unscaled_tensors.append(tensor * inv_scale.item()) return multi_tensor_l2norm_torch(chunk_size, noop_flag, [unscaled_tensors], per_tensor) - def multi_tensor_adam(self, *args, **kwargs): - if not args and not kwargs: - return multi_tensor_adam_torch - return multi_tensor_adam_torch(*args, **kwargs) - - def multi_tensor_adam_param_remainder(self, *args, **kwargs): - if not args and not kwargs: - return multi_tensor_adam_param_remainder_torch - return multi_tensor_adam_param_remainder_torch(*args, **kwargs) - - def multi_tensor_adam_fp8(self, *args, **kwargs) -> None: - raise NotImplementedError("multi_tensor_adam_fp8 - not implemented in reference backend") - - def multi_tensor_adam_capturable(self, *args, **kwargs) -> None: - raise NotImplementedError("multi_tensor_adam_capturable - not implemented in reference backend") - - def multi_tensor_adam_capturable_master(self, *args, **kwargs) -> None: - raise NotImplementedError("multi_tensor_adam_capturable_master - not implemented in reference backend") - - def multi_tensor_sgd(self, *args, **kwargs) -> None: - return multi_tensor_sgd_torch(*args, **kwargs) - - def multi_tensor_compute_scale_and_scale_inv(self, *args, **kwargs) -> None: - raise NotImplementedError("multi_tensor_compute_scale_and_scale_inv - not implemented in reference backend") - - def bulk_overlap_ag_with_external_gemm(self, *args, **kwargs) -> Any: - raise NotImplementedError("bulk_overlap_ag_with_external_gemm - not implemented in reference backend") - - def create_fp8_tensor_meta(self) -> FP8TensorMeta: - return FP8TensorMeta() + def multi_tensor_adam( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: + return multi_tensor_adam_torch( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, step, mode, bias_correction, weight_decay + ) - def create_comm_overlap_helper(self, *args, **kwargs) -> Any: - raise NotImplementedError("create_comm_overlap_helper - not implemented in reference backend") + def multi_tensor_adam_param_remainder( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: + return multi_tensor_adam_param_remainder_torch( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, step, mode, bias_correction, weight_decay + ) - def create_comm_overlap(self, *args, **kwargs) -> Any: - raise NotImplementedError("create_comm_overlap - not implemented in reference backend") + def multi_tensor_sgd( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, + ) -> None: + return multi_tensor_sgd_torch( + chunk_size, noop_flag, tensor_lists, + wd, momentum, dampening, lr, nesterov, first_run, wd_after_momentum, scale + ) - def create_comm_overlap_p2p(self, *args, **kwargs) -> Any: - raise NotImplementedError("create_comm_overlap_p2p - not implemented in reference backend") + def get_flash_attention_class(self): + from .flash_attention import FlashAttentionTorch + return FlashAttentionTorch diff --git a/transformer_engine/plugin/core/backends/reference/register_ops.py b/transformer_engine/plugin/core/backends/reference/register_ops.py index 3d311a6c75..9ecbf10974 100644 --- a/transformer_engine/plugin/core/backends/reference/register_ops.py +++ b/transformer_engine/plugin/core/backends/reference/register_ops.py @@ -43,20 +43,11 @@ def register_builtins(registry) -> None: # Normalization OpImpl(op_name="rmsnorm_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.rmsnorm_fwd, is_avail), vendor=None, priority=50), OpImpl(op_name="rmsnorm_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.rmsnorm_bwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="rmsnorm_bwd_add", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.rmsnorm_bwd_add, is_avail), vendor=None, priority=50), OpImpl(op_name="layernorm_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.layernorm_fwd, is_avail), vendor=None, priority=50), OpImpl(op_name="layernorm_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.layernorm_bwd, is_avail), vendor=None, priority=50), # GEMM OpImpl(op_name="generic_gemm", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.generic_gemm, is_avail), vendor=None, priority=50), - OpImpl(op_name="te_general_grouped_gemm", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.te_general_grouped_gemm, is_avail), vendor=None, priority=50), - - # Quantization - OpImpl(op_name="quantize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.quantize, is_avail), vendor=None, priority=50), - OpImpl(op_name="dequantize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.dequantize, is_avail), vendor=None, priority=50), - OpImpl(op_name="bgrad_quantize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.bgrad_quantize, is_avail), vendor=None, priority=50), - OpImpl(op_name="multi_tensor_quantize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_quantize, is_avail), vendor=None, priority=50), - OpImpl(op_name="split_quantize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.split_quantize, is_avail), vendor=None, priority=50), # Activations - Forward OpImpl(op_name="gelu", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.gelu, is_avail), vendor=None, priority=50), @@ -101,94 +92,25 @@ def register_builtins(registry) -> None: OpImpl(op_name="scaled_aligned_causal_masked_softmax_forward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.scaled_aligned_causal_masked_softmax_forward, is_avail), vendor=None, priority=50), OpImpl(op_name="scaled_aligned_causal_masked_softmax_backward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.scaled_aligned_causal_masked_softmax_backward, is_avail), vendor=None, priority=50), - # MOE operations - OpImpl(op_name="moe_permute_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.moe_permute_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="moe_permute_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.moe_permute_bwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="moe_unpermute_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.moe_unpermute_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="moe_unpermute_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.moe_unpermute_bwd, is_avail), vendor=None, priority=50), - - # Fused attention + # Fused attention backend getter OpImpl(op_name="get_fused_attn_backend", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.get_fused_attn_backend, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_attn_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_attn_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_attn_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_attn_bwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fa_prepare_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fa_prepare_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fa_prepare_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fa_prepare_bwd, is_avail), vendor=None, priority=50), - - # KV cache - OpImpl(op_name="copy_to_kv_cache", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.copy_to_kv_cache, is_avail), vendor=None, priority=50), - - # Tensor format conversions - OpImpl(op_name="convert_thd_to_bshd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.convert_thd_to_bshd, is_avail), vendor=None, priority=50), - OpImpl(op_name="convert_bshd_to_thd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.convert_bshd_to_thd, is_avail), vendor=None, priority=50), - - # RoPE (Rotary Position Embedding) - OpImpl(op_name="fused_rope_forward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_rope_forward, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_rope_backward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_rope_backward, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_qkv_rope_forward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_qkv_rope_forward, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_qkv_rope_backward", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_qkv_rope_backward, is_avail), vendor=None, priority=50), - - # TopK and MOE aux loss - OpImpl(op_name="fused_topk_with_score_function_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_topk_with_score_function_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_topk_with_score_function_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_topk_with_score_function_bwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_score_for_moe_aux_loss_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_score_for_moe_aux_loss_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_score_for_moe_aux_loss_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_score_for_moe_aux_loss_bwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_moe_aux_loss_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_moe_aux_loss_fwd, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_moe_aux_loss_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_moe_aux_loss_bwd, is_avail), vendor=None, priority=50), # Dropout OpImpl(op_name="dropout_fwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.dropout_fwd, is_avail), vendor=None, priority=50), OpImpl(op_name="dropout_bwd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.dropout_bwd, is_avail), vendor=None, priority=50), - # FP8 operations - OpImpl(op_name="fp8_transpose", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fp8_transpose, is_avail), vendor=None, priority=50), - OpImpl(op_name="swap_first_dims", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.swap_first_dims, is_avail), vendor=None, priority=50), - OpImpl(op_name="compute_amax", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.compute_amax, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_amax_and_scale_update_after_reduction", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_amax_and_scale_update_after_reduction, is_avail), vendor=None, priority=50), - OpImpl(op_name="fp8_block_scaling_compute_partial_amax", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fp8_block_scaling_compute_partial_amax, is_avail), vendor=None, priority=50), - OpImpl(op_name="fp8_block_scaling_partial_cast", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fp8_block_scaling_partial_cast, is_avail), vendor=None, priority=50), - - # Padding operations - OpImpl(op_name="fused_multi_row_padding", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_multi_row_padding, is_avail), vendor=None, priority=50), - OpImpl(op_name="fused_multi_row_unpadding", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.fused_multi_row_unpadding, is_avail), vendor=None, priority=50), - # Library version getters OpImpl(op_name="get_cublasLt_version", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.get_cublasLt_version, is_avail), vendor=None, priority=50), OpImpl(op_name="get_cudnn_version", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.get_cudnn_version, is_avail), vendor=None, priority=50), OpImpl(op_name="get_num_cublas_streams", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.get_num_cublas_streams, is_avail), vendor=None, priority=50), - # THD (Tensor, Hidden, Dimension) operations - OpImpl(op_name="thd_read_half_tensor", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_read_half_tensor, is_avail), vendor=None, priority=50), - OpImpl(op_name="thd_second_half_lse_correction", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_second_half_lse_correction, is_avail), vendor=None, priority=50), - OpImpl(op_name="thd_read_second_half_lse", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_read_second_half_lse, is_avail), vendor=None, priority=50), - OpImpl(op_name="thd_out_correction", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_out_correction, is_avail), vendor=None, priority=50), - OpImpl(op_name="thd_grad_correction", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_grad_correction, is_avail), vendor=None, priority=50), - OpImpl(op_name="thd_get_partitioned_indices", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.thd_get_partitioned_indices, is_avail), vendor=None, priority=50), - - # NVSHMEM operations - OpImpl(op_name="init_nvshmem_backend", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.init_nvshmem_backend, is_avail), vendor=None, priority=50), - OpImpl(op_name="create_nvshmem_tensor", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.create_nvshmem_tensor, is_avail), vendor=None, priority=50), - OpImpl(op_name="nvshmem_send_on_current_stream", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.nvshmem_send_on_current_stream, is_avail), vendor=None, priority=50), - OpImpl(op_name="nvshmem_wait_on_current_stream", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.nvshmem_wait_on_current_stream, is_avail), vendor=None, priority=50), - OpImpl(op_name="nvshmem_finalize", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.nvshmem_finalize, is_avail), vendor=None, priority=50), - # Multi-tensor optimizer operations OpImpl(op_name="multi_tensor_scale", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_scale, is_avail), vendor=None, priority=50), OpImpl(op_name="multi_tensor_l2norm", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_l2norm, is_avail), vendor=None, priority=50), OpImpl(op_name="multi_tensor_unscale_l2norm", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_unscale_l2norm, is_avail), vendor=None, priority=50), OpImpl(op_name="multi_tensor_adam", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_adam, is_avail), vendor=None, priority=50), OpImpl(op_name="multi_tensor_adam_param_remainder", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_adam_param_remainder, is_avail), vendor=None, priority=50), - OpImpl(op_name="multi_tensor_adam_fp8", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_adam_fp8, is_avail), vendor=None, priority=50), - OpImpl(op_name="multi_tensor_adam_capturable", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_adam_capturable, is_avail), vendor=None, priority=50), - OpImpl(op_name="multi_tensor_adam_capturable_master", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_adam_capturable_master, is_avail), vendor=None, priority=50), OpImpl(op_name="multi_tensor_sgd", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_sgd, is_avail), vendor=None, priority=50), - OpImpl(op_name="multi_tensor_compute_scale_and_scale_inv", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.multi_tensor_compute_scale_and_scale_inv, is_avail), vendor=None, priority=50), - - # Communication overlap operations - OpImpl(op_name="bulk_overlap_ag_with_external_gemm", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.bulk_overlap_ag_with_external_gemm, is_avail), vendor=None, priority=50), - OpImpl(op_name="create_fp8_tensor_meta", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.create_fp8_tensor_meta, is_avail), vendor=None, priority=50), - OpImpl(op_name="create_comm_overlap_helper", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.create_comm_overlap_helper, is_avail), vendor=None, priority=50), - OpImpl(op_name="create_comm_overlap", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.create_comm_overlap, is_avail), vendor=None, priority=50), - OpImpl(op_name="create_comm_overlap_p2p", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.create_comm_overlap_p2p, is_avail), vendor=None, priority=50), # FlashAttention class getter OpImpl(op_name="get_flash_attention_class", impl_id="reference.torch", kind=BackendImplKind.REFERENCE, fn=_bind_is_available(backend.get_flash_attention_class, is_avail), vendor=None, priority=50), diff --git a/transformer_engine/plugin/core/backends/vendor/cuda/cuda.py b/transformer_engine/plugin/core/backends/vendor/cuda/cuda.py index 98ef965811..8be7dd5052 100644 --- a/transformer_engine/plugin/core/backends/vendor/cuda/cuda.py +++ b/transformer_engine/plugin/core/backends/vendor/cuda/cuda.py @@ -1,12 +1,11 @@ # Copyright (c) 2025, BAAI. All rights reserved. # # See LICENSE for license information. - +import os +import sys from typing import Any, Dict, List, Optional, Tuple, Union - import torch - -from ....ops import TEFLBackendBase, FP8TensorMeta +from ....ops import * def _load_cuda_libs(): import ctypes @@ -115,70 +114,6 @@ def _get_tex(): import transformer_engine_torch_nv return transformer_engine_torch_nv -def _torch_dtype_to_te_dtype(torch_dtype, tex_module): - if torch_dtype is None: - return None - - NativeDType = tex_module.DType - if type(torch_dtype).__name__ == 'DType' and type(torch_dtype).__module__ == 'transformer_engine_torch_nv': - return torch_dtype - - if hasattr(torch_dtype, 'name') and hasattr(torch_dtype, 'value'): - from transformer_engine.plugin.core.ops import DType as PyDType - if isinstance(torch_dtype, PyDType): - dtype_name = torch_dtype.name - if hasattr(NativeDType, dtype_name): - return getattr(NativeDType, dtype_name) - - dtype_map = { - torch.float32: NativeDType.kFloat32, - torch.float16: NativeDType.kFloat16, - torch.bfloat16: NativeDType.kBFloat16, - torch.int32: NativeDType.kInt32, - torch.uint8: NativeDType.kByte, - } - - if hasattr(torch, 'float8_e4m3fn'): - dtype_map[torch.float8_e4m3fn] = NativeDType.kFloat8E4M3 - if hasattr(torch, 'float8_e5m2'): - dtype_map[torch.float8_e5m2] = NativeDType.kFloat8E5M2 - - return dtype_map.get(torch_dtype, torch_dtype) - -def _convert_dtype_params(func): - import functools - import inspect - import os - - @functools.wraps(func) - def wrapper(self, *args, **kwargs): - dtype_params = ['otype', 'output_dtype', 'bias_type'] - - from transformer_engine.plugin.core.ops import DType as PyDType - - def needs_conversion(val): - return isinstance(val, torch.dtype) or isinstance(val, PyDType) - - for param_name in dtype_params: - if param_name in kwargs: - value = kwargs[param_name] - if needs_conversion(value): - converted = self._to_te_dtype(value) - kwargs[param_name] = converted - - sig = inspect.signature(func) - param_names = list(sig.parameters.keys())[1:] - - args_list = list(args) - for i, (param_name, arg_value) in enumerate(zip(param_names, args_list)): - if param_name in dtype_params and needs_conversion(arg_value): - converted = self._to_te_dtype(arg_value) - args_list[i] = converted - - return func(self, *args_list, **kwargs) - - return wrapper - class CUDABackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -192,16 +127,9 @@ def _get_tex(self): self._tex = _get_tex() return self._tex - def _to_te_dtype(self, torch_dtype): - return _torch_dtype_to_te_dtype(torch_dtype, self._get_tex()) - def is_available(self) -> bool: return _check_cuda_available() - def get_flash_attention_class(self): - from .flash_attention import FlashAttentionCUDA - return FlashAttentionCUDA - def get_attention_backend(self, attention_params=None): """ CUDA backend uses the default attention backend selection logic. @@ -214,6 +142,7 @@ def get_attention_backend(self, attention_params=None): from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils return dpa_utils._original_get_attention_backend(attention_params) +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def quantize( self, tensor: torch.Tensor, @@ -224,35 +153,34 @@ def quantize( tex = self._get_tex() return tex.quantize(tensor, quantizer, output, noop) - @_convert_dtype_params def dequantize( self, - input: torch.Tensor, - otype: torch.dtype, - ) -> torch.Tensor: + input: Any, + otype: DType, + ) -> Any: tex = self._get_tex() + otype = tex.DType(int(otype)) if otype is not None else None return tex.dequantize(input, otype) def bgrad_quantize( self, input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: tex = self._get_tex() return tex.bgrad_quantize(input, quantizer) - @_convert_dtype_params def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -261,61 +189,53 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: + ) -> List[Any]: tex = self._get_tex() - - if bias_type is None: - bias_type = self._to_te_dtype(torch.bfloat16) - + + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None + output_dtype = tex.DType(int(output_dtype)) if output_dtype is not None else None return tex.generic_gemm( A, transA, B, transB, D, quantizer, output_dtype, bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, accumulate, use_split_accumulator, comm_overlap, comm_type, extra_output, bulk_overlap, alpha, beta ) - - def te_general_grouped_gemm(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.te_general_grouped_gemm(*args, **kwargs) - + # GELU and variants # def gelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.gelu(input, quantizer) - def geglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.geglu(input, quantizer) def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgelu(input, quantizer) - def qgeglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgeglu(input, quantizer) + # ReLU and variants # def relu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.relu(input, quantizer) - def reglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.reglu(input, quantizer) def srelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.srelu(input, quantizer) - def sreglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.sreglu(input, quantizer) - + # SwiGLU and variants # def silu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.silu(input, quantizer) - def swiglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.swiglu(input, quantizer) @@ -328,42 +248,39 @@ def clamped_swiglu( ) -> Any: tex = self._get_tex() return tex.clamped_swiglu(input, quantizer, limit, alpha) - + # Backward of GELU and variants # def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgelu(grad, fwd_input, quantizer) def dgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgeglu(grad, fwd_input, quantizer) - def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgelu(grad, fwd_input, quantizer) def dqgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgeglu(grad, fwd_input, quantizer) - + # Backward of ReLU and variants # def drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.drelu(grad, fwd_input, quantizer) def dreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dreglu(grad, fwd_input, quantizer) - def dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsrelu(grad, fwd_input, quantizer) def dsreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsreglu(grad, fwd_input, quantizer) - + # Backward of SiLU and variants # def dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsilu(grad, fwd_input, quantizer) def dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dswiglu(grad, fwd_input, quantizer) - def clamped_dswiglu( self, grad: torch.Tensor, @@ -374,131 +291,207 @@ def clamped_dswiglu( ) -> Any: tex = self._get_tex() return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha) - - def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + # DBias + DAct fusions # + def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dgelu(grad, fwd_input, quantizer) - - def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsilu(grad, fwd_input, quantizer) - - def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_drelu(grad, fwd_input, quantizer) - - def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dqgelu(grad, fwd_input, quantizer) - - def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsrelu(grad, fwd_input, quantizer) - - @_convert_dtype_params + # Permutation functions + def moe_permute_fwd( + self, + input: torch.Tensor, + dtype: DType, + indices: torch.Tensor, + num_out_tokens: int, + workspace: List[torch.Tensor], + max_expanded_token_num: int, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_fwd(input, dtype,indices,num_out_tokens,workspace,max_expanded_token_num) + def moe_permute_bwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_bwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_fwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_fwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_bwd( + self, + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_bwd(input_bwd,input_fwd,dtype,row_id_map,prob) + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_forward(input, scale) + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_forward(input, mask, scale_factor) + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_forward(input, scale_factor) + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_backward( + output_grads_, softmax_results_, scale_factor + ) + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_forward(input, scale_factor) + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_backward( + output_grad_, softmax_results_, scale_factor + ) + # Other granular functions def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, mu, rsigma = tex.layernorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.layernorm_fwd( input, weight, bias, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - return y, mu, rsigma - def layernorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dgamma, dbeta = tex.layernorm_bwd(dy, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dgamma, dbeta - - @_convert_dtype_params + return tex.layernorm_bwd( + dz, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma + ) def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, y_quant, rsigma = tex.rmsnorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.rmsnorm_fwd( input, weight, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - if y_quant is not None: - y_quant = y_quant.view(*orig_shape) - return y, y_quant, rsigma - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dw = tex.rmsnorm_bwd(dy, x, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dw - - def rmsnorm_bwd_add(self, *args, **kwargs) -> Any: + return tex.rmsnorm_bwd(dz, x, rsigma, gamma, sm_margin, zero_centered_gamma) + def rmsnorm_bwd_add( + self, + dz: torch.Tensor, + x: torch.Tensor, + add: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.rmsnorm_bwd_add(*args, **kwargs) + return tex.rmsnorm_bwd_add(dz, x, add, rsigma, gamma, sm_margin, zero_centered_gamma) def multi_tensor_quantize( self, @@ -507,7 +500,6 @@ def multi_tensor_quantize( ) -> List[Any]: tex = self._get_tex() return tex.multi_tensor_quantize(tensor_list, quantizer_list) - def split_quantize( self, tensor: torch.Tensor, @@ -516,246 +508,457 @@ def split_quantize( ) -> List[Any]: tex = self._get_tex() return tex.split_quantize(tensor, split_sections, quantizer_list) - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_fwd(*args, **kwargs) - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_bwd(*args, **kwargs) - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_fwd(*args, **kwargs) - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_bwd(*args, **kwargs) - - def scaled_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_forward(input, scale) - - def scaled_softmax_backward( + def te_general_grouped_gemm( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_backward(output_grad, softmax_output, scale) - - def scaled_masked_softmax_forward( + A: List[Any], + transa: bool, + B: List[Any], + transb: bool, + D: Optional[List[torch.Tensor]], + D_type: DType, + m_splits: List[int], + bias: List[torch.Tensor], + bias_type: DType, + single_output: bool, + pre_gelu_out: List[torch.Tensor], + grad: bool, + workspace: List[torch.Tensor], + workspaceSizes: int, + accumulate: bool, + use_split_accumulator: bool, + math_sm_count: int, + ) -> Optional[List[torch.Tensor]]: + tex = self._get_tex() + D_type = tex.DType(int(D_type)) if D_type is not None else None + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + return tex.te_general_grouped_gemm( + A, transa, B, transb, D, D_type, m_splits, bias, bias_type, + single_output, pre_gelu_out, grad, workspace, workspaceSizes, + accumulate, use_split_accumulator, math_sm_count + ) + def fp8_transpose( self, input: torch.Tensor, - mask: torch.Tensor, - scale: float, + dtype: DType, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_forward(input, mask, scale) - - def scaled_masked_softmax_backward( + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.fp8_transpose(input, dtype, out) + def swap_first_dims( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + tensor: torch.Tensor, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_backward(output_grad, softmax_output, scale) + return tex.swap_first_dims(tensor, out) + def get_fused_attn_backend( + self, + is_training: bool, + q_dtype: DType, + kv_dtype: DType, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + p_dropout: float, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + window_size_left: int, + window_size_right: int, + return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: + tex = self._get_tex() + + q_dtype = tex.DType(int(q_dtype)) if q_dtype is not None else None + kv_dtype = tex.DType(int(kv_dtype)) if kv_dtype is not None else None + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + result = tex.get_fused_attn_backend( + is_training, q_dtype, kv_dtype, qkv_layout, bias_type, + attn_mask_type, softmax_type, p_dropout, num_attn_heads, + num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, window_size_left, window_size_right, return_max_logit + ) + return NVTE_Fused_Attn_Backend(result) - def scaled_upper_triang_masked_softmax_forward( + def compute_amax( self, input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax: torch.Tensor, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_forward(input, scale) - - def scaled_upper_triang_masked_softmax_backward( + return tex.compute_amax(input, amax) + def fused_amax_and_scale_update_after_reduction( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax_reduction_buffer: torch.Tensor, + amax_histories: List[torch.Tensor], + scales: List[torch.Tensor], + amax_compute_algo: str, + fp8_dtype: DType, + margin: float, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_backward(output_grad, softmax_output, scale) - - def scaled_aligned_causal_masked_softmax_forward( + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.fused_amax_and_scale_update_after_reduction( + amax_reduction_buffer, amax_histories, scales, + amax_compute_algo, fp8_dtype, margin + ) + def fp8_block_scaling_compute_partial_amax( self, - input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_forward(input, scale) - - def scaled_aligned_causal_masked_softmax_backward( + return tex.fp8_block_scaling_compute_partial_amax( + tensor, amax, h, w, start_offset, block_len + ) + def fp8_block_scaling_partial_cast( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + inp: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + out_dtype: DType, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_backward(output_grad, softmax_output, scale) - - def get_fused_attn_backend(self, *args, **kwargs) -> int: + out_dtype = tex.DType(int(out_dtype)) if out_dtype is not None else None + return tex.fp8_block_scaling_partial_cast( + inp, out, scale, h, w, start_offset, block_len, out_dtype + ) + def fused_multi_row_padding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + padded_input_row_list: List[int], + ) -> None: tex = self._get_tex() - - args_list = list(args) - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - - if hasattr(py_enum, 'value'): - enum_value = int(py_enum.value) - for member_name in dir(native_enum_class): - if not member_name.startswith('_'): - try: - member = getattr(native_enum_class, member_name) - if hasattr(member, 'value') and int(member.value) == enum_value: - return member - except: - pass - - if hasattr(py_enum, 'value'): - return int(py_enum.value) - - return py_enum - - if len(args) > 1: - args_list[1] = self._to_te_dtype(args[1]) - if len(args) > 2: - args_list[2] = self._to_te_dtype(args[2]) - if len(args) > 3: - args_list[3] = convert_enum(args[3], tex.NVTE_QKV_Layout) - if len(args) > 4: - args_list[4] = convert_enum(args[4], tex.NVTE_Bias_Type) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_Mask_Type) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Softmax_Type) - - return tex.get_fused_attn_backend(*args_list, **kwargs) - - def fused_attn_fwd(self, *args, **kwargs) -> Any: + return tex.fused_multi_row_padding( + input, output, input_row_list, padded_input_row_list + ) + def fused_multi_row_unpadding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + unpadded_input_row_list: List[int], + ) -> None: tex = self._get_tex() + return tex.fused_multi_row_unpadding( + input, output, input_row_list, unpadded_input_row_list + ) - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_QKV_Layout) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Bias_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Mask_Type) - if len(args) > 9: - args_list[9] = convert_enum(args[9], tex.NVTE_Softmax_Type) - - return tex.fused_attn_fwd(*args_list, **kwargs) - - def fused_attn_bwd(self, *args, **kwargs) -> Any: + # attention kernels + def fa_prepare_fwd( + self, + qkvi: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_QKV_Layout) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Bias_Type) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Mask_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Softmax_Type) - if len(args) > 19: - args_list[19] = self._to_te_dtype(args[19]) - - if 'dqkv_dtype' in kwargs: - kwargs['dqkv_dtype'] = self._to_te_dtype(kwargs['dqkv_dtype']) - - return tex.fused_attn_bwd(*args_list, **kwargs) - - def fa_prepare_fwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_fwd(qkvi) + def fa_prepare_bwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fa_prepare_fwd(*args, **kwargs) - - def fa_prepare_bwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_bwd(q, k, v) + def fused_attn_fwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + is_training: bool, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + fake_dtype: torch.dtype, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + s_quantizer: Any, + o_quantizer: Any, + Bias: Optional[torch.Tensor], + SoftmaxOffset: Optional[torch.Tensor], + rng_gen: Optional[torch.Generator], + rng_elts_per_thread: int, + return_max_logit: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.fa_prepare_bwd(*args, **kwargs) - def copy_to_kv_cache(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + return tex.fused_attn_fwd( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + Bias, + SoftmaxOffset, + rng_gen, + rng_elts_per_thread, + return_max_logit + ) + def fused_attn_bwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + O: Any, + dO: Any, + fake_dtype: torch.dtype, + dqkv_type: DType, + Aux_CTX_Tensors: List[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + s_quantizer: Any, + dp_quantizer: Any, + dqkv_quantizer: Any, + ) -> List[Any]: tex = self._get_tex() - return tex.copy_to_kv_cache(*args, **kwargs) - def convert_thd_to_bshd(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + dqkv_type = tex.DType(int(dqkv_type)) if dqkv_type is not None else None + + return tex.fused_attn_bwd( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + O, + dO, + fake_dtype, + dqkv_type, + Aux_CTX_Tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + s_quantizer, + dp_quantizer, + dqkv_quantizer + ) + def copy_to_kv_cache( + self, + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: NVTE_QKV_Format, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, + ) -> None: tex = self._get_tex() - return tex.convert_thd_to_bshd(*args, **kwargs) - - def convert_bshd_to_thd(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.copy_to_kv_cache( + new_k, + new_v, + k_cache, + v_cache, + page_table, + cu_new_lens, + cu_cached_lens, + qkv_format, + b, + max_ctx_len, + max_seq_len, + max_pages_per_seq, + is_non_paged + ) + def convert_thd_to_bshd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + b: int, + max_seq_len: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.convert_bshd_to_thd(*args, **kwargs) - - def fused_rope_forward(self, *args, **kwargs) -> Any: + return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len) + def convert_bshd_to_thd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + t: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_rope_forward(*args, **kwargs) + return tex.convert_bshd_to_thd(tensor, cu_seqlens, t) - def fused_rope_backward(self, *args, **kwargs) -> Any: + # 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: tex = self._get_tex() - return tex.fused_rope_backward(*args, **kwargs) - - def fused_qkv_rope_forward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_forward( + 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, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_qkv_rope_forward(*args, **kwargs) - - def fused_qkv_rope_backward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_backward( + output_grads, freqs, 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]: tex = self._get_tex() - return tex.fused_qkv_rope_backward(*args, **kwargs) + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_qkv_rope_forward( + 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: + tex = self._get_tex() + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.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 + ) + # fused router def fused_topk_with_score_function_fwd( self, logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, - score_function: Any, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, expert_bias: Optional[torch.Tensor], - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_topk_with_score_function_fwd( - logits, topk, use_pre_softmax, num_groups, group_topk, - scaling_factor, score_function, expert_bias + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, ) - def fused_topk_with_score_function_bwd( self, num_tokens: int, @@ -765,24 +968,33 @@ def fused_topk_with_score_function_bwd( grad_probs: torch.Tensor, topk: int, use_pre_softmax: bool, - scaling_factor: float, - score_function: Any, - ) -> Any: + scaling_factor: Optional[float], + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_topk_with_score_function_bwd( - num_tokens, num_experts, routing_map, intermediate_output, - grad_probs, topk, use_pre_softmax, scaling_factor, score_function + num_tokens, + num_experts, + routing_map, + intermediate_output, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, ) - def fused_score_for_moe_aux_loss_fwd( self, logits: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.fused_score_for_moe_aux_loss_fwd(logits, topk, score_function) - + return tex.fused_score_for_moe_aux_loss_fwd( + logits, + topk, + score_function, + ) def fused_score_for_moe_aux_loss_bwd( self, num_tokens: int, @@ -790,13 +1002,17 @@ def fused_score_for_moe_aux_loss_bwd( intermediate_output: torch.Tensor, grad_scores: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_score_for_moe_aux_loss_bwd( - num_tokens, num_experts, intermediate_output, grad_scores, topk, score_function + num_tokens, + num_experts, + intermediate_output, + grad_scores, + topk, + score_function, ) - def fused_moe_aux_loss_fwd( self, probs: torch.Tensor, @@ -807,13 +1023,18 @@ def fused_moe_aux_loss_fwd( num_cols: int, topk: int, coeff: float, - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_moe_aux_loss_fwd( - probs, tokens_per_expert, total_num_tokens, num_experts, - num_rows, num_cols, topk, coeff + probs, + tokens_per_expert, + total_num_tokens, + num_experts, + num_rows, + num_cols, + topk, + coeff, ) - def fused_moe_aux_loss_bwd( self, Const_buf: torch.Tensor, @@ -821,152 +1042,146 @@ def fused_moe_aux_loss_bwd( num_rows: int, num_cols: int, grad_aux_loss: torch.Tensor, - ) -> Any: + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_moe_aux_loss_bwd( - Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss - ) + return tex.fused_moe_aux_loss_bwd(Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss) + # Dropout def dropout_fwd( self, input: torch.Tensor, dropout_probability: float, - out: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.dropout_fwd(input, dropout_probability, out) - def dropout_bwd( self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, - grad_input: Optional[torch.Tensor] = None, + grad_input: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() return tex.dropout_bwd(grad_output, mask, dropout_probability, grad_input) - def fp8_transpose( - self, - input: torch.Tensor, - dtype: Any, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.fp8_transpose(input, dtype, out=out) - - def swap_first_dims( - self, - tensor: torch.Tensor, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.swap_first_dims(tensor, out=out) - - def compute_amax( - self, - input: torch.Tensor, - amax: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.compute_amax(input, amax) - - def fused_amax_and_scale_update_after_reduction(self, *args, **kwargs) -> None: - tex = self._get_tex() - tex.fused_amax_and_scale_update_after_reduction(*args, **kwargs) - - def fp8_block_scaling_compute_partial_amax( - self, - tensor: torch.Tensor, - amax: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_compute_partial_amax(tensor, amax, h, w, start_offset, block_len) - - def fp8_block_scaling_partial_cast( - self, - inp: torch.Tensor, - out: torch.Tensor, - scale: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - out_dtype: Any, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_partial_cast(inp, out, scale, h, w, start_offset, block_len, out_dtype) - - def fused_multi_row_padding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_padding(*args, **kwargs) - - def fused_multi_row_unpadding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_unpadding(*args, **kwargs) - + # Misc def get_cublasLt_version(self) -> int: tex = self._get_tex() return tex.get_cublasLt_version() - def get_cudnn_version(self) -> int: tex = self._get_tex() return tex.get_cudnn_version() - def get_num_cublas_streams(self) -> int: tex = self._get_tex() return tex.get_num_cublas_streams() - def thd_read_half_tensor(self, *args, **kwargs) -> Any: + # Support THD format for Context Parallel + def thd_read_half_tensor( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + half_idx: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_half_tensor(*args, **kwargs) - - def thd_second_half_lse_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_half_tensor(tensor, cu_seqlens, half_idx) + def thd_second_half_lse_correction( + self, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_second_half_lse_correction(*args, **kwargs) - - def thd_read_second_half_lse(self, *args, **kwargs) -> Any: + return tex.thd_second_half_lse_correction( + lse, lse_per_step, cu_seqlens, lse_packed + ) + def thd_read_second_half_lse( + self, + lse: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + second_half_lse_seqlen: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_second_half_lse(*args, **kwargs) - - def thd_out_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_second_half_lse( + lse, cu_seqlens, lse_packed, second_half_lse_seqlen + ) + def thd_out_correction( + self, + out: torch.Tensor, + out_per_step: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + only_second_half: bool, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_out_correction(*args, **kwargs) - - def thd_grad_correction(self, *args, **kwargs) -> Any: + return tex.thd_out_correction( + out, out_per_step, lse, lse_per_step, + cu_seqlens, only_second_half, lse_packed + ) + def thd_grad_correction( + self, + grad: torch.Tensor, + grad_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + first_half: str, + second_half: str, + ) -> None: tex = self._get_tex() - return tex.thd_grad_correction(*args, **kwargs) - - def thd_get_partitioned_indices(self, *args, **kwargs) -> Any: + return tex.thd_grad_correction( + grad, grad_per_step, cu_seqlens, + first_half, second_half + ) + def thd_get_partitioned_indices( + self, + cu_seqlens: torch.Tensor, + total_tokens: int, + world_size: int, + rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_get_partitioned_indices(*args, **kwargs) + return tex.thd_get_partitioned_indices( + cu_seqlens, total_tokens, world_size, rank + ) - def init_nvshmem_backend(self, *args, **kwargs) -> None: + # nvshmem functions + def init_nvshmem_backend( + self, + process_group: Any, + ) -> None: tex = self._get_tex() - tex.init_nvshmem_backend(*args, **kwargs) - - def create_nvshmem_tensor(self, *args, **kwargs) -> torch.Tensor: + return tex.init_nvshmem_backend(process_group) + def create_nvshmem_tensor( + self, + shape: List[int], + dtype: torch.dtype, + ) -> torch.Tensor: tex = self._get_tex() - return tex.create_nvshmem_tensor(*args, **kwargs) - - def nvshmem_send_on_current_stream(self, *args, **kwargs) -> None: + return tex.create_nvshmem_tensor(shape, dtype) + def nvshmem_send_on_current_stream( + self, + src: torch.Tensor, + dst: torch.Tensor, + peer: int, + signal: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.nvshmem_send_on_current_stream(*args, **kwargs) - - def nvshmem_wait_on_current_stream(self, *args, **kwargs) -> None: + return tex.nvshmem_send_on_current_stream(src, dst, peer, signal) + def nvshmem_wait_on_current_stream( + self, + signal: torch.Tensor, + wait_kind: str, + ) -> None: tex = self._get_tex() - tex.nvshmem_wait_on_current_stream(*args, **kwargs) - + return tex.nvshmem_wait_on_current_stream(signal, wait_kind) def nvshmem_finalize(self) -> None: tex = self._get_tex() - tex.nvshmem_finalize() + return tex.nvshmem_finalize() + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -975,98 +1190,195 @@ def multi_tensor_scale( scale: float, ) -> None: tex = self._get_tex() - tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) - + return tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.multi_tensor_l2norm(chunk_size, noop_flag, tensor_lists, per_tensor) - def multi_tensor_unscale_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - scale: torch.Tensor, - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, scale, per_tensor) - + return tex.multi_tensor_unscale_l2norm( + chunk_size, noop_flag, tensor_lists, inv_scale, per_tensor + ) def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): - tex = self._get_tex() - if chunk_size is None: - return tex.multi_tensor_adam - tex.multi_tensor_adam( - chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, - eps, step, mode, bias_correction, weight_decay + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: + tex = self._get_tex() + return tex.multi_tensor_adam( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay ) - - def multi_tensor_adam_param_remainder(self, *args, **kwargs) -> None: + def multi_tensor_adam_param_remainder( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_param_remainder(*args, **kwargs) - - def multi_tensor_adam_fp8(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_param_remainder( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay + ) + def multi_tensor_adam_fp8( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + fp8_dtype: DType, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_fp8(*args, **kwargs) - - def multi_tensor_adam_capturable(self, *args, **kwargs) -> None: + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.multi_tensor_adam_fp8( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + fp8_dtype + ) + def multi_tensor_adam_capturable( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable(*args, **kwargs) - - def multi_tensor_adam_capturable_master(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_adam_capturable_master( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable_master(*args, **kwargs) - - def multi_tensor_sgd(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable_master( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_sgd( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_sgd(*args, **kwargs) - - def multi_tensor_compute_scale_and_scale_inv(self, *args, **kwargs) -> None: + return tex.multi_tensor_sgd( + chunk_size, noop_flag, tensor_lists, + wd, momentum, dampening, + lr, nesterov, first_run, + wd_after_momentum, scale + ) + def multi_tensor_compute_scale_and_scale_inv( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + max_fp8: float, + force_pow_2_scales: bool, + epsilon: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_compute_scale_and_scale_inv(*args, **kwargs) + return tex.multi_tensor_compute_scale_and_scale_inv( + chunk_size, noop_flag, tensor_lists, + max_fp8, force_pow_2_scales, epsilon + ) + # Comm+GEMM Overlap def bulk_overlap_ag_with_external_gemm( self, - allgather_communicator: Any, + allgather_communicator: CommOverlap, send_stream: Any, recv_stream: Any, ) -> Any: tex = self._get_tex() return tex.bulk_overlap_ag_with_external_gemm(allgather_communicator, send_stream, recv_stream) +############## class func ################################# + def get_flash_attention_class(self): + from .flash_attention import FlashAttentionCUDA + return FlashAttentionCUDA def create_fp8_tensor_meta(self) -> FP8TensorMeta: tex = self._get_tex() return tex.FP8TensorMeta() - def create_comm_overlap_helper( self, world_group: Optional[Any] = None, intra_node_group: Optional[Any] = None, - ) -> Any: + ) -> "CommOverlapHelper": tex = self._get_tex() - if world_group is None: - return tex.CommOverlapHelper() return tex.CommOverlapHelper(world_group, intra_node_group) - def create_comm_overlap( self, buffer_shape: List[int], @@ -1082,7 +1394,7 @@ def create_comm_overlap( set_sm_margin: bool = True, atomic_gemm: bool = False, rs_overlap_first_gemm: bool = False, - ) -> Any: + ) -> "CommOverlap": tex = self._get_tex() return tex.CommOverlap( buffer_shape, buffer_dtype, helper, tp_size, @@ -1090,7 +1402,6 @@ def create_comm_overlap( gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm ) - def create_comm_overlap_p2p( self, buffer_shape: List[int], @@ -1107,7 +1418,7 @@ def create_comm_overlap_p2p( atomic_gemm: bool = False, use_ce: bool = True, aggregate: bool = False, - ) -> Any: + ) -> "CommOverlapP2P": tex = self._get_tex() return tex.CommOverlapP2P( buffer_shape, buffer_dtype, helper, tp_size, comm_type, diff --git a/transformer_engine/plugin/core/backends/vendor/hygon/hygon.py b/transformer_engine/plugin/core/backends/vendor/hygon/hygon.py index 92e8868ed9..c87aef8430 100644 --- a/transformer_engine/plugin/core/backends/vendor/hygon/hygon.py +++ b/transformer_engine/plugin/core/backends/vendor/hygon/hygon.py @@ -5,10 +5,8 @@ import os import sys from typing import Any, Dict, List, Optional, Tuple, Union - import torch - -from ....ops import TEFLBackendBase, FP8TensorMeta, NVTE_Fused_Attn_Backend +from ....ops import * def _load_hygon_libs(): import ctypes @@ -78,69 +76,6 @@ def _get_tex(): import transformer_engine_torch_hygon return transformer_engine_torch_hygon -def _torch_dtype_to_te_dtype(torch_dtype, tex_module): - if torch_dtype is None: - return None - - NativeDType = tex_module.DType - if type(torch_dtype).__name__ == 'DType' and type(torch_dtype).__module__ == 'transformer_engine_torch_hygon': - return torch_dtype - - if hasattr(torch_dtype, 'name') and hasattr(torch_dtype, 'value'): - from transformer_engine.plugin.core.ops import DType as PyDType - if isinstance(torch_dtype, PyDType): - dtype_name = torch_dtype.name - if hasattr(NativeDType, dtype_name): - return getattr(NativeDType, dtype_name) - - dtype_map = { - torch.float32: NativeDType.kFloat32, - torch.float16: NativeDType.kFloat16, - torch.bfloat16: NativeDType.kBFloat16, - torch.int32: NativeDType.kInt32, - torch.uint8: NativeDType.kByte, - } - - if hasattr(torch, 'float8_e4m3fn'): - dtype_map[torch.float8_e4m3fn] = NativeDType.kFloat8E4M3 - if hasattr(torch, 'float8_e5m2'): - dtype_map[torch.float8_e5m2] = NativeDType.kFloat8E5M2 - - return dtype_map.get(torch_dtype, torch_dtype) - -def _convert_dtype_params(func): - import functools - import inspect - - @functools.wraps(func) - def wrapper(self, *args, **kwargs): - dtype_params = ['otype', 'output_dtype', 'bias_type'] - - from transformer_engine.plugin.core.ops import DType as PyDType - - def needs_conversion(val): - return isinstance(val, torch.dtype) or isinstance(val, PyDType) - - for param_name in dtype_params: - if param_name in kwargs: - value = kwargs[param_name] - if needs_conversion(value): - converted = self._to_te_dtype(value) - kwargs[param_name] = converted - - sig = inspect.signature(func) - param_names = list(sig.parameters.keys())[1:] - - args_list = list(args) - for i, (param_name, arg_value) in enumerate(zip(param_names, args_list)): - if param_name in dtype_params and needs_conversion(arg_value): - converted = self._to_te_dtype(arg_value) - args_list[i] = converted - - return func(self, *args_list, **kwargs) - - return wrapper - class HygonBackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -154,16 +89,9 @@ def _get_tex(self): self._tex = _get_tex() return self._tex - def _to_te_dtype(self, torch_dtype): - return _torch_dtype_to_te_dtype(torch_dtype, self._get_tex()) - def is_available(self) -> bool: return _check_hygon_available() - def get_flash_attention_class(self): - from .flash_attention import FlashAttentionHYGON - return FlashAttentionHYGON - def get_attention_backend(self, attention_params=None): from packaging.version import Version as PkgVersion from ....logger_manager import get_logger @@ -196,6 +124,7 @@ def get_attention_backend(self, attention_params=None): available_backends, ) +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def quantize( self, tensor: torch.Tensor, @@ -206,35 +135,34 @@ def quantize( tex = self._get_tex() return tex.quantize(tensor, quantizer, output, noop) - @_convert_dtype_params def dequantize( self, - input: torch.Tensor, - otype: torch.dtype, - ) -> torch.Tensor: + input: Any, + otype: DType, + ) -> Any: tex = self._get_tex() + otype = tex.DType(int(otype)) if otype is not None else None return tex.dequantize(input, otype) def bgrad_quantize( self, input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: tex = self._get_tex() return tex.bgrad_quantize(input, quantizer) - @_convert_dtype_params def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -243,68 +171,56 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: + ) -> List[Any]: tex = self._get_tex() - - if bias_type is None: - bias_type = self._to_te_dtype(torch.bfloat16) - + + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None + output_dtype = tex.DType(int(output_dtype)) if output_dtype is not None else None return tex.generic_gemm( A, transA, B, transB, D, quantizer, output_dtype, bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, accumulate, use_split_accumulator, comm_overlap, comm_type, extra_output, bulk_overlap, alpha, beta ) - - def te_general_grouped_gemm(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.te_general_grouped_gemm(*args, **kwargs) - + # GELU and variants # def gelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.gelu(input, quantizer) - def geglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.geglu(input, quantizer) - def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgelu(input, quantizer) - def qgeglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgeglu(input, quantizer) - + # ReLU and variants # def relu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.relu(input, quantizer) - def reglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.reglu(input, quantizer) - def srelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.srelu(input, quantizer) - def sreglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.sreglu(input, quantizer) - + # SwiGLU and variants # def silu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.silu(input, quantizer) - def swiglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.swiglu(input, quantizer) - def clamped_swiglu( self, input: torch.Tensor, @@ -314,47 +230,39 @@ def clamped_swiglu( ) -> Any: tex = self._get_tex() return tex.clamped_swiglu(input, quantizer, limit, alpha) - + # Backward of GELU and variants # def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgelu(grad, fwd_input, quantizer) - def dgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgeglu(grad, fwd_input, quantizer) - def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgelu(grad, fwd_input, quantizer) - def dqgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgeglu(grad, fwd_input, quantizer) - + # Backward of ReLU and variants # def drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.drelu(grad, fwd_input, quantizer) - def dreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dreglu(grad, fwd_input, quantizer) - def dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsrelu(grad, fwd_input, quantizer) - def dsreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsreglu(grad, fwd_input, quantizer) - + # Backward of SiLU and variants # def dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsilu(grad, fwd_input, quantizer) - def dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dswiglu(grad, fwd_input, quantizer) - def clamped_dswiglu( self, grad: torch.Tensor, @@ -365,131 +273,207 @@ def clamped_dswiglu( ) -> Any: tex = self._get_tex() return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha) - - def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + # DBias + DAct fusions # + def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dgelu(grad, fwd_input, quantizer) - - def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsilu(grad, fwd_input, quantizer) - - def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_drelu(grad, fwd_input, quantizer) - - def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dqgelu(grad, fwd_input, quantizer) - - def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsrelu(grad, fwd_input, quantizer) - - @_convert_dtype_params + # Permutation functions + def moe_permute_fwd( + self, + input: torch.Tensor, + dtype: DType, + indices: torch.Tensor, + num_out_tokens: int, + workspace: List[torch.Tensor], + max_expanded_token_num: int, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_fwd(input, dtype,indices,num_out_tokens,workspace,max_expanded_token_num) + def moe_permute_bwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_bwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_fwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_fwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_bwd( + self, + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_bwd(input_bwd,input_fwd,dtype,row_id_map,prob) + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_forward(input, scale) + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_forward(input, mask, scale_factor) + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_forward(input, scale_factor) + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_backward( + output_grads_, softmax_results_, scale_factor + ) + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_forward(input, scale_factor) + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_backward( + output_grad_, softmax_results_, scale_factor + ) + # Other granular functions def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, mu, rsigma = tex.layernorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.layernorm_fwd( input, weight, bias, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - return y, mu, rsigma - def layernorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dgamma, dbeta = tex.layernorm_bwd(dy, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dgamma, dbeta - - @_convert_dtype_params + return tex.layernorm_bwd( + dz, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma + ) def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, y_quant, rsigma = tex.rmsnorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.rmsnorm_fwd( input, weight, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - if y_quant is not None: - y_quant = y_quant.view(*orig_shape) - return y, y_quant, rsigma - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dw = tex.rmsnorm_bwd(dy, x, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dw - - def rmsnorm_bwd_add(self, *args, **kwargs) -> Any: + return tex.rmsnorm_bwd(dz, x, rsigma, gamma, sm_margin, zero_centered_gamma) + def rmsnorm_bwd_add( + self, + dz: torch.Tensor, + x: torch.Tensor, + add: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.rmsnorm_bwd_add(*args, **kwargs) + return tex.rmsnorm_bwd_add(dz, x, add, rsigma, gamma, sm_margin, zero_centered_gamma) def multi_tensor_quantize( self, @@ -498,7 +482,6 @@ def multi_tensor_quantize( ) -> List[Any]: tex = self._get_tex() return tex.multi_tensor_quantize(tensor_list, quantizer_list) - def split_quantize( self, tensor: torch.Tensor, @@ -507,150 +490,457 @@ def split_quantize( ) -> List[Any]: tex = self._get_tex() return tex.split_quantize(tensor, split_sections, quantizer_list) - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_fwd(*args, **kwargs) - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_bwd(*args, **kwargs) - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_fwd(*args, **kwargs) - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_bwd(*args, **kwargs) - - def scaled_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: + def te_general_grouped_gemm( + self, + A: List[Any], + transa: bool, + B: List[Any], + transb: bool, + D: Optional[List[torch.Tensor]], + D_type: DType, + m_splits: List[int], + bias: List[torch.Tensor], + bias_type: DType, + single_output: bool, + pre_gelu_out: List[torch.Tensor], + grad: bool, + workspace: List[torch.Tensor], + workspaceSizes: int, + accumulate: bool, + use_split_accumulator: bool, + math_sm_count: int, + ) -> Optional[List[torch.Tensor]]: + tex = self._get_tex() + D_type = tex.DType(int(D_type)) if D_type is not None else None + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + return tex.te_general_grouped_gemm( + A, transa, B, transb, D, D_type, m_splits, bias, bias_type, + single_output, pre_gelu_out, grad, workspace, workspaceSizes, + accumulate, use_split_accumulator, math_sm_count + ) + def fp8_transpose( + self, + input: torch.Tensor, + dtype: DType, + out: Optional[torch.Tensor], + ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_softmax_forward(input, scale) - - def scaled_softmax_backward( + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.fp8_transpose(input, dtype, out) + def swap_first_dims( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + tensor: torch.Tensor, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_softmax_backward(output_grad, softmax_output, scale) + return tex.swap_first_dims(tensor, out) + def get_fused_attn_backend( + self, + is_training: bool, + q_dtype: DType, + kv_dtype: DType, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + p_dropout: float, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + window_size_left: int, + window_size_right: int, + return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: + tex = self._get_tex() + + q_dtype = tex.DType(int(q_dtype)) if q_dtype is not None else None + kv_dtype = tex.DType(int(kv_dtype)) if kv_dtype is not None else None + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + result = tex.get_fused_attn_backend( + is_training, q_dtype, kv_dtype, qkv_layout, bias_type, + attn_mask_type, softmax_type, p_dropout, num_attn_heads, + num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, window_size_left, window_size_right, return_max_logit + ) + return NVTE_Fused_Attn_Backend(result) - def scaled_masked_softmax_forward( + def compute_amax( self, input: torch.Tensor, - mask: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax: torch.Tensor, + ) -> None: tex = self._get_tex() - return tex.scaled_masked_softmax_forward(input, mask, scale) - - def scaled_masked_softmax_backward( + return tex.compute_amax(input, amax) + def fused_amax_and_scale_update_after_reduction( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax_reduction_buffer: torch.Tensor, + amax_histories: List[torch.Tensor], + scales: List[torch.Tensor], + amax_compute_algo: str, + fp8_dtype: DType, + margin: float, + ) -> None: tex = self._get_tex() - return tex.scaled_masked_softmax_backward(output_grad, softmax_output, scale) - - def scaled_upper_triang_masked_softmax_forward( + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.fused_amax_and_scale_update_after_reduction( + amax_reduction_buffer, amax_histories, scales, + amax_compute_algo, fp8_dtype, margin + ) + def fp8_block_scaling_compute_partial_amax( + self, + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + ) -> None: + tex = self._get_tex() + return tex.fp8_block_scaling_compute_partial_amax( + tensor, amax, h, w, start_offset, block_len + ) + def fp8_block_scaling_partial_cast( + self, + inp: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + out_dtype: DType, + ) -> None: + tex = self._get_tex() + out_dtype = tex.DType(int(out_dtype)) if out_dtype is not None else None + return tex.fp8_block_scaling_partial_cast( + inp, out, scale, h, w, start_offset, block_len, out_dtype + ) + def fused_multi_row_padding( self, input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + output: torch.Tensor, + input_row_list: List[int], + padded_input_row_list: List[int], + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_forward(input, scale) - - def scaled_upper_triang_masked_softmax_backward( + return tex.fused_multi_row_padding( + input, output, input_row_list, padded_input_row_list + ) + def fused_multi_row_unpadding( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + unpadded_input_row_list: List[int], + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_backward(output_grad, softmax_output, scale) + return tex.fused_multi_row_unpadding( + input, output, input_row_list, unpadded_input_row_list + ) - def scaled_aligned_causal_masked_softmax_forward( + # attention kernels + def fa_prepare_fwd( self, - input: torch.Tensor, - scale: float, + qkvi: torch.Tensor, ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_forward(input, scale) - - def scaled_aligned_causal_masked_softmax_backward( + return tex.fa_prepare_fwd(qkvi) + def fa_prepare_bwd( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_backward(output_grad, softmax_output, scale) - - def get_fused_attn_backend(self, *args, **kwargs) -> int: - raise NotImplementedError("get_fused_attn_backend - not implemented in hygon backend") - - def fused_attn_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_attn_fwd - not implemented in hygon backend") - - def fused_attn_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError("fused_attn_bwd - not implemented in hygon backend") - - def fa_prepare_fwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_bwd(q, k, v) + def fused_attn_fwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + is_training: bool, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + fake_dtype: torch.dtype, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + s_quantizer: Any, + o_quantizer: Any, + Bias: Optional[torch.Tensor], + SoftmaxOffset: Optional[torch.Tensor], + rng_gen: Optional[torch.Generator], + rng_elts_per_thread: int, + return_max_logit: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.fa_prepare_fwd(*args, **kwargs) - def fa_prepare_bwd(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + return tex.fused_attn_fwd( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + Bias, + SoftmaxOffset, + rng_gen, + rng_elts_per_thread, + return_max_logit + ) + def fused_attn_bwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + O: Any, + dO: Any, + fake_dtype: torch.dtype, + dqkv_type: DType, + Aux_CTX_Tensors: List[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + s_quantizer: Any, + dp_quantizer: Any, + dqkv_quantizer: Any, + ) -> List[Any]: tex = self._get_tex() - return tex.fa_prepare_bwd(*args, **kwargs) - def copy_to_kv_cache(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + dqkv_type = tex.DType(int(dqkv_type)) if dqkv_type is not None else None + + return tex.fused_attn_bwd( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + O, + dO, + fake_dtype, + dqkv_type, + Aux_CTX_Tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + s_quantizer, + dp_quantizer, + dqkv_quantizer + ) + def copy_to_kv_cache( + self, + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: NVTE_QKV_Format, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, + ) -> None: tex = self._get_tex() - return tex.copy_to_kv_cache(*args, **kwargs) - - def convert_thd_to_bshd(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.copy_to_kv_cache( + new_k, + new_v, + k_cache, + v_cache, + page_table, + cu_new_lens, + cu_cached_lens, + qkv_format, + b, + max_ctx_len, + max_seq_len, + max_pages_per_seq, + is_non_paged + ) + def convert_thd_to_bshd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + b: int, + max_seq_len: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.convert_thd_to_bshd(*args, **kwargs) - - def convert_bshd_to_thd(self, *args, **kwargs) -> Any: + return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len) + def convert_bshd_to_thd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + t: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.convert_bshd_to_thd(*args, **kwargs) + return tex.convert_bshd_to_thd(tensor, cu_seqlens, t) - def fused_rope_forward(self, *args, **kwargs) -> Any: + # 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: tex = self._get_tex() - return tex.fused_rope_forward(*args, **kwargs) - - def fused_rope_backward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_forward( + 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, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_rope_backward(*args, **kwargs) - - def fused_qkv_rope_forward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_backward( + output_grads, freqs, 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]: tex = self._get_tex() - return tex.fused_qkv_rope_forward(*args, **kwargs) - - def fused_qkv_rope_backward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_qkv_rope_forward( + 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: tex = self._get_tex() - return tex.fused_qkv_rope_backward(*args, **kwargs) + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.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 + ) + # fused router def fused_topk_with_score_function_fwd( self, logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, - score_function: Any, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, expert_bias: Optional[torch.Tensor], - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_topk_with_score_function_fwd( - logits, topk, use_pre_softmax, num_groups, group_topk, - scaling_factor, score_function, expert_bias + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, ) - def fused_topk_with_score_function_bwd( self, num_tokens: int, @@ -660,24 +950,33 @@ def fused_topk_with_score_function_bwd( grad_probs: torch.Tensor, topk: int, use_pre_softmax: bool, - scaling_factor: float, - score_function: Any, - ) -> Any: + scaling_factor: Optional[float], + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_topk_with_score_function_bwd( - num_tokens, num_experts, routing_map, intermediate_output, - grad_probs, topk, use_pre_softmax, scaling_factor, score_function + num_tokens, + num_experts, + routing_map, + intermediate_output, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, ) - def fused_score_for_moe_aux_loss_fwd( self, logits: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.fused_score_for_moe_aux_loss_fwd(logits, topk, score_function) - + return tex.fused_score_for_moe_aux_loss_fwd( + logits, + topk, + score_function, + ) def fused_score_for_moe_aux_loss_bwd( self, num_tokens: int, @@ -685,13 +984,17 @@ def fused_score_for_moe_aux_loss_bwd( intermediate_output: torch.Tensor, grad_scores: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_score_for_moe_aux_loss_bwd( - num_tokens, num_experts, intermediate_output, grad_scores, topk, score_function + num_tokens, + num_experts, + intermediate_output, + grad_scores, + topk, + score_function, ) - def fused_moe_aux_loss_fwd( self, probs: torch.Tensor, @@ -702,13 +1005,18 @@ def fused_moe_aux_loss_fwd( num_cols: int, topk: int, coeff: float, - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_moe_aux_loss_fwd( - probs, tokens_per_expert, total_num_tokens, num_experts, - num_rows, num_cols, topk, coeff + probs, + tokens_per_expert, + total_num_tokens, + num_experts, + num_rows, + num_cols, + topk, + coeff, ) - def fused_moe_aux_loss_bwd( self, Const_buf: torch.Tensor, @@ -716,147 +1024,146 @@ def fused_moe_aux_loss_bwd( num_rows: int, num_cols: int, grad_aux_loss: torch.Tensor, - ) -> Any: + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_moe_aux_loss_bwd( - Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss - ) + return tex.fused_moe_aux_loss_bwd(Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss) + # Dropout def dropout_fwd( self, input: torch.Tensor, dropout_probability: float, - out: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.dropout_fwd(input, dropout_probability, out) - def dropout_bwd( self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, - grad_input: Optional[torch.Tensor] = None, + grad_input: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() return tex.dropout_bwd(grad_output, mask, dropout_probability, grad_input) - def fp8_transpose( - self, - input: torch.Tensor, - dtype: Any, - *, - out: torch.Tensor, - ) -> None: + # Misc + def get_cublasLt_version(self) -> int: tex = self._get_tex() - tex.fp8_transpose(input, dtype, out=out) + return tex.get_cublasLt_version() + def get_cudnn_version(self) -> int: + tex = self._get_tex() + return tex.get_cudnn_version() + def get_num_cublas_streams(self) -> int: + tex = self._get_tex() + return tex.get_num_cublas_streams() - def swap_first_dims( + # Support THD format for Context Parallel + def thd_read_half_tensor( self, tensor: torch.Tensor, - *, - out: torch.Tensor, - ) -> None: + cu_seqlens: torch.Tensor, + half_idx: int, + ) -> torch.Tensor: tex = self._get_tex() - tex.swap_first_dims(tensor, out=out) - - def compute_amax( + return tex.thd_read_half_tensor(tensor, cu_seqlens, half_idx) + def thd_second_half_lse_correction( self, - input: torch.Tensor, - amax: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, ) -> None: tex = self._get_tex() - tex.compute_amax(input, amax) - - def fused_amax_and_scale_update_after_reduction(self, *args, **kwargs) -> None: - tex = self._get_tex() - tex.fused_amax_and_scale_update_after_reduction(*args, **kwargs) - - def fp8_block_scaling_compute_partial_amax( + return tex.thd_second_half_lse_correction( + lse, lse_per_step, cu_seqlens, lse_packed + ) + def thd_read_second_half_lse( self, - tensor: torch.Tensor, - amax: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - ) -> None: + lse: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + second_half_lse_seqlen: int, + ) -> torch.Tensor: tex = self._get_tex() - tex.fp8_block_scaling_compute_partial_amax(tensor, amax, h, w, start_offset, block_len) - - def fp8_block_scaling_partial_cast( + return tex.thd_read_second_half_lse( + lse, cu_seqlens, lse_packed, second_half_lse_seqlen + ) + def thd_out_correction( self, - inp: torch.Tensor, out: torch.Tensor, - scale: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - out_dtype: Any, + out_per_step: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + only_second_half: bool, + lse_packed: bool, ) -> None: tex = self._get_tex() - tex.fp8_block_scaling_partial_cast(inp, out, scale, h, w, start_offset, block_len, out_dtype) - - def fused_multi_row_padding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_padding(*args, **kwargs) - - def fused_multi_row_unpadding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_unpadding(*args, **kwargs) - - def get_cublasLt_version(self) -> int: - tex = self._get_tex() - return tex.get_cublasLt_version() - - def get_cudnn_version(self) -> int: - tex = self._get_tex() - return tex.get_cudnn_version() - - def get_num_cublas_streams(self) -> int: - tex = self._get_tex() - return tex.get_num_cublas_streams() - - def thd_read_half_tensor(self, *args, **kwargs) -> Any: + return tex.thd_out_correction( + out, out_per_step, lse, lse_per_step, + cu_seqlens, only_second_half, lse_packed + ) + def thd_grad_correction( + self, + grad: torch.Tensor, + grad_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + first_half: str, + second_half: str, + ) -> None: tex = self._get_tex() - return tex.thd_read_half_tensor(*args, **kwargs) - - def thd_second_half_lse_correction(self, *args, **kwargs) -> Any: + return tex.thd_grad_correction( + grad, grad_per_step, cu_seqlens, + first_half, second_half + ) + def thd_get_partitioned_indices( + self, + cu_seqlens: torch.Tensor, + total_tokens: int, + world_size: int, + rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_second_half_lse_correction(*args, **kwargs) + return tex.thd_get_partitioned_indices( + cu_seqlens, total_tokens, world_size, rank + ) - def thd_read_second_half_lse(self, *args, **kwargs) -> Any: + # nvshmem functions + def init_nvshmem_backend( + self, + process_group: Any, + ) -> None: tex = self._get_tex() - return tex.thd_read_second_half_lse(*args, **kwargs) - - def thd_out_correction(self, *args, **kwargs) -> Any: + return tex.init_nvshmem_backend(process_group) + def create_nvshmem_tensor( + self, + shape: List[int], + dtype: torch.dtype, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_out_correction(*args, **kwargs) - - def thd_grad_correction(self, *args, **kwargs) -> Any: + return tex.create_nvshmem_tensor(shape, dtype) + def nvshmem_send_on_current_stream( + self, + src: torch.Tensor, + dst: torch.Tensor, + peer: int, + signal: torch.Tensor, + ) -> None: tex = self._get_tex() - return tex.thd_grad_correction(*args, **kwargs) - - def thd_get_partitioned_indices(self, *args, **kwargs) -> Any: + return tex.nvshmem_send_on_current_stream(src, dst, peer, signal) + def nvshmem_wait_on_current_stream( + self, + signal: torch.Tensor, + wait_kind: str, + ) -> None: tex = self._get_tex() - return tex.thd_get_partitioned_indices(*args, **kwargs) - - def init_nvshmem_backend(self, *args, **kwargs) -> None: - raise NotImplementedError("init_nvshmem_backend - not implemented in hygon backend") - - def create_nvshmem_tensor(self, *args, **kwargs) -> torch.Tensor: - raise NotImplementedError("create_nvshmem_tensor - not implemented in hygon backend") - - def nvshmem_send_on_current_stream(self, *args, **kwargs) -> None: - raise NotImplementedError("nvshmem_send_on_current_stream - not implemented in hygon backend") - - def nvshmem_wait_on_current_stream(self, *args, **kwargs) -> None: - raise NotImplementedError("nvshmem_wait_on_current_stream - not implemented in hygon backend") - + return tex.nvshmem_wait_on_current_stream(signal, wait_kind) def nvshmem_finalize(self) -> None: - raise NotImplementedError("nvshmem_finalize - not implemented in hygon backend") + tex = self._get_tex() + return tex.nvshmem_finalize() + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -865,98 +1172,195 @@ def multi_tensor_scale( scale: float, ) -> None: tex = self._get_tex() - tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) - + return tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.multi_tensor_l2norm(chunk_size, noop_flag, tensor_lists, per_tensor) - def multi_tensor_unscale_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - scale: torch.Tensor, - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, scale, per_tensor) - + return tex.multi_tensor_unscale_l2norm( + chunk_size, noop_flag, tensor_lists, inv_scale, per_tensor + ) def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: tex = self._get_tex() - if chunk_size is None: - return tex.multi_tensor_adam - tex.multi_tensor_adam( - chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, - eps, step, mode, bias_correction, weight_decay + return tex.multi_tensor_adam( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay ) - - def multi_tensor_adam_param_remainder(self, *args, **kwargs) -> None: + def multi_tensor_adam_param_remainder( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_param_remainder(*args, **kwargs) - - def multi_tensor_adam_fp8(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_param_remainder( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay + ) + def multi_tensor_adam_fp8( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + fp8_dtype: DType, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_fp8(*args, **kwargs) - - def multi_tensor_adam_capturable(self, *args, **kwargs) -> None: + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.multi_tensor_adam_fp8( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + fp8_dtype + ) + def multi_tensor_adam_capturable( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable(*args, **kwargs) - - def multi_tensor_adam_capturable_master(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_adam_capturable_master( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable_master(*args, **kwargs) - - def multi_tensor_sgd(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable_master( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_sgd( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_sgd(*args, **kwargs) - - def multi_tensor_compute_scale_and_scale_inv(self, *args, **kwargs) -> None: + return tex.multi_tensor_sgd( + chunk_size, noop_flag, tensor_lists, + wd, momentum, dampening, + lr, nesterov, first_run, + wd_after_momentum, scale + ) + def multi_tensor_compute_scale_and_scale_inv( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + max_fp8: float, + force_pow_2_scales: bool, + epsilon: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_compute_scale_and_scale_inv(*args, **kwargs) + return tex.multi_tensor_compute_scale_and_scale_inv( + chunk_size, noop_flag, tensor_lists, + max_fp8, force_pow_2_scales, epsilon + ) + # Comm+GEMM Overlap def bulk_overlap_ag_with_external_gemm( self, - allgather_communicator: Any, + allgather_communicator: CommOverlap, send_stream: Any, recv_stream: Any, ) -> Any: tex = self._get_tex() return tex.bulk_overlap_ag_with_external_gemm(allgather_communicator, send_stream, recv_stream) +############## class func ################################# + def get_flash_attention_class(self): + from .flash_attention import FlashAttentionHYGON + return FlashAttentionHYGON def create_fp8_tensor_meta(self) -> FP8TensorMeta: tex = self._get_tex() return tex.FP8TensorMeta() - def create_comm_overlap_helper( self, world_group: Optional[Any] = None, intra_node_group: Optional[Any] = None, - ) -> Any: + ) -> "CommOverlapHelper": tex = self._get_tex() - if world_group is None: - return tex.CommOverlapHelper() return tex.CommOverlapHelper(world_group, intra_node_group) - def create_comm_overlap( self, buffer_shape: List[int], @@ -972,7 +1376,7 @@ def create_comm_overlap( set_sm_margin: bool = True, atomic_gemm: bool = False, rs_overlap_first_gemm: bool = False, - ) -> Any: + ) -> "CommOverlap": tex = self._get_tex() return tex.CommOverlap( buffer_shape, buffer_dtype, helper, tp_size, @@ -980,7 +1384,6 @@ def create_comm_overlap( gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm ) - def create_comm_overlap_p2p( self, buffer_shape: List[int], @@ -997,7 +1400,7 @@ def create_comm_overlap_p2p( atomic_gemm: bool = False, use_ce: bool = True, aggregate: bool = False, - ) -> Any: + ) -> "CommOverlapP2P": tex = self._get_tex() return tex.CommOverlapP2P( buffer_shape, buffer_dtype, helper, tp_size, comm_type, diff --git a/transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py b/transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py index 5013fa7c23..294e79fcb9 100644 --- a/transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py +++ b/transformer_engine/plugin/core/backends/vendor/iluvatar/iluvatar.py @@ -7,8 +7,7 @@ import math import torch -from ....ops import TEFLBackendBase, FP8TensorMeta - +from ....ops import * def _load_iluvatar_libs(): import ctypes @@ -105,67 +104,6 @@ def _get_tex(): import transformer_engine_iluvatar.pytorch.ixte_torch return transformer_engine_iluvatar.pytorch.ixte_torch -def _torch_dtype_to_te_dtype(torch_dtype, tex_module): - if torch_dtype is None: - return None - - NativeDType = tex_module.DType - if type(torch_dtype).__name__ == 'DType' and type(torch_dtype).__module__ == 'transformer_engine_iluvatar.pytorch.ixte_torch': - return torch_dtype - - if hasattr(torch_dtype, 'name') and hasattr(torch_dtype, 'value'): - from transformer_engine.plugin.core.ops import DType as PyDType - if isinstance(torch_dtype, PyDType): - dtype_name = torch_dtype.name - if hasattr(NativeDType, dtype_name): - return getattr(NativeDType, dtype_name) - - dtype_map = { - torch.uint8: NativeDType.kByte, - torch.float8_e4m3fn: NativeDType.kFloat8E4M3, - torch.float8_e5m2: NativeDType.kFloat8E5M2, - torch.int32: NativeDType.kInt32, - torch.float32: NativeDType.kFloat32, - torch.half: NativeDType.kFloat16, - torch.bfloat16: NativeDType.kBFloat16, - } - - return dtype_map.get(torch_dtype, torch_dtype) - -def _convert_dtype_params(func): - import functools - import inspect - import os - - @functools.wraps(func) - def wrapper(self, *args, **kwargs): - dtype_params = ['otype', 'output_dtype', 'bias_type'] - - from transformer_engine.plugin.core.ops import DType as PyDType - - def needs_conversion(val): - return isinstance(val, torch.dtype) or isinstance(val, PyDType) - - for param_name in dtype_params: - if param_name in kwargs: - value = kwargs[param_name] - if needs_conversion(value): - converted = self._to_te_dtype(value) - kwargs[param_name] = converted - - sig = inspect.signature(func) - param_names = list(sig.parameters.keys())[1:] - - args_list = list(args) - for i, (param_name, arg_value) in enumerate(zip(param_names, args_list)): - if param_name in dtype_params and needs_conversion(arg_value): - converted = self._to_te_dtype(arg_value) - args_list[i] = converted - - return func(self, *args_list, **kwargs) - - return wrapper - class IluvatarBackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -179,18 +117,42 @@ def _get_tex(self): self._tex = _get_tex() return self._tex - def _to_te_dtype(self, torch_dtype): - return _torch_dtype_to_te_dtype(torch_dtype, self._get_tex()) - def is_available(self) -> bool: return _check_iluvatar_available() - - def get_flash_attention_class(self): - raise NotImplementedError("get_flash_attention_class - not implemented in iluvatar backend") def get_attention_backend(self, attention_params=None): - raise NotImplementedError("get_attention_backend - not implemented in iluvatar backend") - + from packaging.version import Version as PkgVersion + from ....logger_manager import get_logger + logger = get_logger() + + # Read environment variables to determine which backends to enable + use_flash_attention = int(os.getenv("NVTE_FLASH_ATTN", "1")) + use_fused_attention = int(os.getenv("NVTE_FUSED_ATTN", "1")) + use_unfused_attention = int(os.getenv("NVTE_UNFUSED_ATTN", "1")) + + # Log disabled backends + if not use_flash_attention: + logger.info_once("Disabling FlashAttention due to NVTE_FLASH_ATTN=0") + if not use_fused_attention: + logger.info_once("Disabling FusedAttention due to NVTE_FUSED_ATTN=0") + if not use_unfused_attention: + logger.info_once("Disabling UnfusedDotProductAttention due to NVTE_UNFUSED_ATTN=0") + + flash_attention_backend = PkgVersion("2.6.0") if use_flash_attention else None + fused_attention_backend = NVTE_Fused_Attn_Backend.NVTE_No_Backend + + available_backends = [use_flash_attention, use_fused_attention, use_unfused_attention] + + return ( + use_flash_attention, + flash_attention_backend, + use_fused_attention, + fused_attention_backend, + use_unfused_attention, + available_backends, + ) + +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def quantize( self, tensor: torch.Tensor, @@ -201,35 +163,34 @@ def quantize( tex = self._get_tex() return tex.quantize(tensor, quantizer, output, noop) - @_convert_dtype_params def dequantize( self, - input: torch.Tensor, - otype: torch.dtype, - ) -> torch.Tensor: + input: Any, + otype: DType, + ) -> Any: tex = self._get_tex() + otype = tex.DType(int(otype)) if otype is not None else None return tex.dequantize(input, otype) def bgrad_quantize( self, input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: tex = self._get_tex() return tex.bgrad_quantize(input, quantizer) - @_convert_dtype_params def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -238,119 +199,98 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: - # Check shape + ) -> List[Any]: tex = self._get_tex() - - if bias_type is None: - bias_type = self._to_te_dtype(torch.bfloat16) - + + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None + output_dtype = tex.DType(int(output_dtype)) if output_dtype is not None else None return tex.generic_gemm( A, transA, B, transB, D, quantizer, output_dtype, bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, accumulate, use_split_accumulator, comm_overlap, comm_type, extra_output, bulk_overlap, alpha, beta ) - - def te_general_grouped_gemm(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.te_general_grouped_gemm(*args, **kwargs) - + # GELU and variants # def gelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.gelu(input, quantizer) - def geglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.geglu(input, quantizer) - def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgelu(input, quantizer) - def qgeglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgeglu(input, quantizer) - + # ReLU and variants # def relu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.relu(input, quantizer) - def reglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.reglu(input, quantizer) - def srelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.srelu(input, quantizer) - def sreglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.sreglu(input, quantizer) - + # SwiGLU and variants # def silu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.silu(input, quantizer) - def swiglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.swiglu(input, quantizer) - def clamped_swiglu( - self, - input: torch.Tensor, - quantizer: Any, - limit: float = 7.0, - alpha: float = 1.702, - ) -> Any: + self, + input: torch.Tensor, + quantizer: Any, + limit: float = 7.0, + alpha: float = 1.702, + ) -> Any: tex = self._get_tex() return tex.clamped_swiglu(input, quantizer, limit, alpha) - + # Backward of GELU and variants # def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgelu(grad, fwd_input, quantizer) - def dgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgeglu(grad, fwd_input, quantizer) - def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgelu(grad, fwd_input, quantizer) - def dqgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgeglu(grad, fwd_input, quantizer) - + # Backward of ReLU and variants # def drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.drelu(grad, fwd_input, quantizer) - def dreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dreglu(grad, fwd_input, quantizer) - def dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsrelu(grad, fwd_input, quantizer) - def dsreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsreglu(grad, fwd_input, quantizer) - + # Backward of SiLU and variants # def dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsilu(grad, fwd_input, quantizer) - def dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dswiglu(grad, fwd_input, quantizer) - def clamped_dswiglu( self, grad: torch.Tensor, @@ -361,131 +301,207 @@ def clamped_dswiglu( ) -> Any: tex = self._get_tex() return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha) - - def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + # DBias + DAct fusions # + def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dgelu(grad, fwd_input, quantizer) - - def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsilu(grad, fwd_input, quantizer) - - def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_drelu(grad, fwd_input, quantizer) - - def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dqgelu(grad, fwd_input, quantizer) - - def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsrelu(grad, fwd_input, quantizer) - - @_convert_dtype_params + # Permutation functions + def moe_permute_fwd( + self, + input: torch.Tensor, + dtype: DType, + indices: torch.Tensor, + num_out_tokens: int, + workspace: List[torch.Tensor], + max_expanded_token_num: int, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_fwd(input, dtype,indices,num_out_tokens,workspace,max_expanded_token_num) + def moe_permute_bwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_bwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_fwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_fwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_bwd( + self, + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_bwd(input_bwd,input_fwd,dtype,row_id_map,prob) + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_forward(input, scale) + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_forward(input, mask, scale_factor) + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_forward(input, scale_factor) + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_backward( + output_grads_, softmax_results_, scale_factor + ) + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_forward(input, scale_factor) + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_backward( + output_grad_, softmax_results_, scale_factor + ) + # Other granular functions def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, mu, rsigma = tex.layernorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.layernorm_fwd( input, weight, bias, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - return y, mu, rsigma - def layernorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dgamma, dbeta = tex.layernorm_bwd(dy, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dgamma, dbeta - - @_convert_dtype_params + return tex.layernorm_bwd( + dz, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma + ) def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, y_quant, rsigma = tex.rmsnorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.rmsnorm_fwd( input, weight, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - if y_quant is not None: - y_quant = y_quant.view(*orig_shape) - return y, y_quant, rsigma - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dw = tex.rmsnorm_bwd(dy, x, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dw - - def rmsnorm_bwd_add(self, *args, **kwargs) -> Any: + return tex.rmsnorm_bwd(dz, x, rsigma, gamma, sm_margin, zero_centered_gamma) + def rmsnorm_bwd_add( + self, + dz: torch.Tensor, + x: torch.Tensor, + add: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.rmsnorm_bwd_add(*args, **kwargs) + return tex.rmsnorm_bwd_add(dz, x, add, rsigma, gamma, sm_margin, zero_centered_gamma) def multi_tensor_quantize( self, @@ -494,7 +510,6 @@ def multi_tensor_quantize( ) -> List[Any]: tex = self._get_tex() return tex.multi_tensor_quantize(tensor_list, quantizer_list) - def split_quantize( self, tensor: torch.Tensor, @@ -503,249 +518,457 @@ def split_quantize( ) -> List[Any]: tex = self._get_tex() return tex.split_quantize(tensor, split_sections, quantizer_list) - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex._moe_permute_fwd(*args, **kwargs) - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex._moe_permute_bwd(*args, **kwargs) - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex._moe_unpermute_fwd(*args, **kwargs) - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex._moe_unpermute_bwd(*args, **kwargs) - - def scaled_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_forward(input, scale) - - def scaled_softmax_backward( + def te_general_grouped_gemm( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_backward(output_grad, softmax_output, scale) - - def scaled_masked_softmax_forward( + A: List[Any], + transa: bool, + B: List[Any], + transb: bool, + D: Optional[List[torch.Tensor]], + D_type: DType, + m_splits: List[int], + bias: List[torch.Tensor], + bias_type: DType, + single_output: bool, + pre_gelu_out: List[torch.Tensor], + grad: bool, + workspace: List[torch.Tensor], + workspaceSizes: int, + accumulate: bool, + use_split_accumulator: bool, + math_sm_count: int, + ) -> Optional[List[torch.Tensor]]: + tex = self._get_tex() + D_type = tex.DType(int(D_type)) if D_type is not None else None + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + return tex.te_general_grouped_gemm( + A, transa, B, transb, D, D_type, m_splits, bias, bias_type, + single_output, pre_gelu_out, grad, workspace, workspaceSizes, + accumulate, use_split_accumulator, math_sm_count + ) + def fp8_transpose( self, input: torch.Tensor, - mask: torch.Tensor, - scale: float, + dtype: DType, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_forward(input, mask, scale) - - def scaled_masked_softmax_backward( + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.fp8_transpose(input, dtype, out) + def swap_first_dims( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + tensor: torch.Tensor, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_backward(output_grad, softmax_output, scale) + return tex.swap_first_dims(tensor, out) + def get_fused_attn_backend( + self, + is_training: bool, + q_dtype: DType, + kv_dtype: DType, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + p_dropout: float, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + window_size_left: int, + window_size_right: int, + return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: + tex = self._get_tex() + + q_dtype = tex.DType(int(q_dtype)) if q_dtype is not None else None + kv_dtype = tex.DType(int(kv_dtype)) if kv_dtype is not None else None + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + result = tex.get_fused_attn_backend( + is_training, q_dtype, kv_dtype, qkv_layout, bias_type, + attn_mask_type, softmax_type, p_dropout, num_attn_heads, + num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, window_size_left, window_size_right, return_max_logit + ) + return NVTE_Fused_Attn_Backend(result) - def scaled_upper_triang_masked_softmax_forward( + def compute_amax( self, input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax: torch.Tensor, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_forward(input, scale) - - def scaled_upper_triang_masked_softmax_backward( + return tex.compute_amax(input, amax) + def fused_amax_and_scale_update_after_reduction( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax_reduction_buffer: torch.Tensor, + amax_histories: List[torch.Tensor], + scales: List[torch.Tensor], + amax_compute_algo: str, + fp8_dtype: DType, + margin: float, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_backward(output_grad, softmax_output, scale) - - def scaled_aligned_causal_masked_softmax_forward( + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.fused_amax_and_scale_update_after_reduction( + amax_reduction_buffer, amax_histories, scales, + amax_compute_algo, fp8_dtype, margin + ) + def fp8_block_scaling_compute_partial_amax( self, - input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_forward(input, scale) - - def scaled_aligned_causal_masked_softmax_backward( + return tex.fp8_block_scaling_compute_partial_amax( + tensor, amax, h, w, start_offset, block_len + ) + def fp8_block_scaling_partial_cast( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + inp: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + out_dtype: DType, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_backward(output_grad, softmax_output, scale) - - def get_fused_attn_backend(self, *args, **kwargs) -> int: + out_dtype = tex.DType(int(out_dtype)) if out_dtype is not None else None + return tex.fp8_block_scaling_partial_cast( + inp, out, scale, h, w, start_offset, block_len, out_dtype + ) + def fused_multi_row_padding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + padded_input_row_list: List[int], + ) -> None: tex = self._get_tex() - - args_list = list(args) - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - - if hasattr(py_enum, 'value'): - enum_value = int(py_enum.value) - for member_name in dir(native_enum_class): - if not member_name.startswith('_'): - try: - member = getattr(native_enum_class, member_name) - if hasattr(member, 'value') and int(member.value) == enum_value: - return member - except: - pass - - if hasattr(py_enum, 'value'): - return int(py_enum.value) - - return py_enum - - if len(args) > 1: - args_list[1] = self._to_te_dtype(args[1]) - if len(args) > 2: - args_list[2] = self._to_te_dtype(args[2]) - if len(args) > 3: - args_list[3] = convert_enum(args[3], tex.NVTE_QKV_Layout) - if len(args) > 4: - args_list[4] = convert_enum(args[4], tex.NVTE_Bias_Type) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_Mask_Type) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Softmax_Type) - - return tex.get_fused_attn_backend(*args_list, **kwargs) - - def fused_attn_fwd(self, *args, **kwargs) -> Any: + return tex.fused_multi_row_padding( + input, output, input_row_list, padded_input_row_list + ) + def fused_multi_row_unpadding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + unpadded_input_row_list: List[int], + ) -> None: tex = self._get_tex() + return tex.fused_multi_row_unpadding( + input, output, input_row_list, unpadded_input_row_list + ) - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_QKV_Layout) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Bias_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Mask_Type) - if len(args) > 9: - args_list[9] = convert_enum(args[9], tex.NVTE_Softmax_Type) - - return tex.fused_attn_fwd(*args_list, **kwargs) - - def fused_attn_bwd(self, *args, **kwargs) -> Any: + # attention kernels + def fa_prepare_fwd( + self, + qkvi: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_nv': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_QKV_Layout) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Bias_Type) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Mask_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Softmax_Type) - if len(args) > 19: - args_list[19] = self._to_te_dtype(args[19]) - - if 'dqkv_dtype' in kwargs: - kwargs['dqkv_dtype'] = self._to_te_dtype(kwargs['dqkv_dtype']) - - return tex.fused_attn_bwd(*args_list, **kwargs) - - def fa_prepare_fwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_fwd(qkvi) + def fa_prepare_bwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fa_prepare_fwd(*args, **kwargs) - - def fa_prepare_bwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_bwd(q, k, v) + def fused_attn_fwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + is_training: bool, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + fake_dtype: torch.dtype, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + s_quantizer: Any, + o_quantizer: Any, + Bias: Optional[torch.Tensor], + SoftmaxOffset: Optional[torch.Tensor], + rng_gen: Optional[torch.Generator], + rng_elts_per_thread: int, + return_max_logit: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.fa_prepare_bwd(*args, **kwargs) - def copy_to_kv_cache(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + return tex.fused_attn_fwd( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + Bias, + SoftmaxOffset, + rng_gen, + rng_elts_per_thread, + return_max_logit + ) + def fused_attn_bwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + O: Any, + dO: Any, + fake_dtype: torch.dtype, + dqkv_type: DType, + Aux_CTX_Tensors: List[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + s_quantizer: Any, + dp_quantizer: Any, + dqkv_quantizer: Any, + ) -> List[Any]: tex = self._get_tex() - return tex.copy_to_kv_cache(*args, **kwargs) - def convert_thd_to_bshd(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + dqkv_type = tex.DType(int(dqkv_type)) if dqkv_type is not None else None + + return tex.fused_attn_bwd( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + O, + dO, + fake_dtype, + dqkv_type, + Aux_CTX_Tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + s_quantizer, + dp_quantizer, + dqkv_quantizer + ) + def copy_to_kv_cache( + self, + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: NVTE_QKV_Format, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, + ) -> None: tex = self._get_tex() - return tex.convert_thd_to_bshd(*args, **kwargs) - - def convert_bshd_to_thd(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.copy_to_kv_cache( + new_k, + new_v, + k_cache, + v_cache, + page_table, + cu_new_lens, + cu_cached_lens, + qkv_format, + b, + max_ctx_len, + max_seq_len, + max_pages_per_seq, + is_non_paged + ) + def convert_thd_to_bshd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + b: int, + max_seq_len: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.convert_bshd_to_thd(*args, **kwargs) - - def fused_rope_forward(self, *args, **kwargs) -> Any: - assert args[2] is None, "[Iluvatar] fused_rope_forward does not support start_position now." - assert args[3].name == "NVTE_SBHD", f"[Iluvatar] fused_rope_forward expect NVTE_SBHD, but got {args[3].name}." + return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len) + def convert_bshd_to_thd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + t: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_rope_forward(args[0], args[1], False, False, 1.0) + return tex.convert_bshd_to_thd(tensor, cu_seqlens, t) - def fused_rope_backward(self, *args, **kwargs) -> Any: - assert args[2].name == "NVTE_SBHD", f"[Iluvatar] fused_rope_backward expect NVTE_SBHD, but got {args[2].name}." + # 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: tex = self._get_tex() - return tex.fused_rope_backward(args[0], args[1], False, False, 1.0) - - def fused_qkv_rope_forward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_forward( + 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, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_qkv_rope_forward(*args, **kwargs) - - def fused_qkv_rope_backward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_backward( + output_grads, freqs, 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]: + tex = self._get_tex() + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_qkv_rope_forward( + 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: tex = self._get_tex() - return tex.fused_qkv_rope_backward(*args, **kwargs) + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.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 + ) + # fused router def fused_topk_with_score_function_fwd( self, logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, - score_function: Any, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, expert_bias: Optional[torch.Tensor], - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_topk_with_score_function_fwd( - logits, topk, use_pre_softmax, num_groups, group_topk, - scaling_factor, score_function, expert_bias + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, ) - def fused_topk_with_score_function_bwd( self, num_tokens: int, @@ -755,24 +978,33 @@ def fused_topk_with_score_function_bwd( grad_probs: torch.Tensor, topk: int, use_pre_softmax: bool, - scaling_factor: float, - score_function: Any, - ) -> Any: + scaling_factor: Optional[float], + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_topk_with_score_function_bwd( - num_tokens, num_experts, routing_map, intermediate_output, - grad_probs, topk, use_pre_softmax, scaling_factor, score_function + num_tokens, + num_experts, + routing_map, + intermediate_output, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, ) - def fused_score_for_moe_aux_loss_fwd( self, logits: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.fused_score_for_moe_aux_loss_fwd(logits, topk, score_function) - + return tex.fused_score_for_moe_aux_loss_fwd( + logits, + topk, + score_function, + ) def fused_score_for_moe_aux_loss_bwd( self, num_tokens: int, @@ -780,13 +1012,17 @@ def fused_score_for_moe_aux_loss_bwd( intermediate_output: torch.Tensor, grad_scores: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_score_for_moe_aux_loss_bwd( - num_tokens, num_experts, intermediate_output, grad_scores, topk, score_function + num_tokens, + num_experts, + intermediate_output, + grad_scores, + topk, + score_function, ) - def fused_moe_aux_loss_fwd( self, probs: torch.Tensor, @@ -797,13 +1033,18 @@ def fused_moe_aux_loss_fwd( num_cols: int, topk: int, coeff: float, - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_moe_aux_loss_fwd( - probs, tokens_per_expert, total_num_tokens, num_experts, - num_rows, num_cols, topk, coeff + probs, + tokens_per_expert, + total_num_tokens, + num_experts, + num_rows, + num_cols, + topk, + coeff, ) - def fused_moe_aux_loss_bwd( self, Const_buf: torch.Tensor, @@ -811,152 +1052,146 @@ def fused_moe_aux_loss_bwd( num_rows: int, num_cols: int, grad_aux_loss: torch.Tensor, - ) -> Any: + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_moe_aux_loss_bwd( - Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss - ) + return tex.fused_moe_aux_loss_bwd(Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss) + # Dropout def dropout_fwd( self, input: torch.Tensor, dropout_probability: float, - out: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.dropout_fwd(input, dropout_probability, out) - def dropout_bwd( self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, - grad_input: Optional[torch.Tensor] = None, + grad_input: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() return tex.dropout_bwd(grad_output, mask, dropout_probability, grad_input) - def fp8_transpose( - self, - input: torch.Tensor, - dtype: Any, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.fp8_transpose(input, dtype, out=out) - - def swap_first_dims( - self, - tensor: torch.Tensor, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.swap_first_dims(tensor, out=out) - - def compute_amax( - self, - input: torch.Tensor, - amax: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.compute_amax(input, amax) - - def fused_amax_and_scale_update_after_reduction(self, *args, **kwargs) -> None: - tex = self._get_tex() - tex.fused_amax_and_scale_update_after_reduction(*args, **kwargs) - - def fp8_block_scaling_compute_partial_amax( - self, - tensor: torch.Tensor, - amax: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_compute_partial_amax(tensor, amax, h, w, start_offset, block_len) - - def fp8_block_scaling_partial_cast( - self, - inp: torch.Tensor, - out: torch.Tensor, - scale: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - out_dtype: Any, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_partial_cast(inp, out, scale, h, w, start_offset, block_len, out_dtype) - - def fused_multi_row_padding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_padding(*args, **kwargs) - - def fused_multi_row_unpadding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_unpadding(*args, **kwargs) - + # Misc def get_cublasLt_version(self) -> int: tex = self._get_tex() return tex.get_cublasLt_version() - def get_cudnn_version(self) -> int: tex = self._get_tex() return tex.get_cudnn_version() - def get_num_cublas_streams(self) -> int: tex = self._get_tex() return tex.get_num_cublas_streams() - def thd_read_half_tensor(self, *args, **kwargs) -> Any: + # Support THD format for Context Parallel + def thd_read_half_tensor( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + half_idx: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_half_tensor(*args, **kwargs) - - def thd_second_half_lse_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_half_tensor(tensor, cu_seqlens, half_idx) + def thd_second_half_lse_correction( + self, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_second_half_lse_correction(*args, **kwargs) - - def thd_read_second_half_lse(self, *args, **kwargs) -> Any: + return tex.thd_second_half_lse_correction( + lse, lse_per_step, cu_seqlens, lse_packed + ) + def thd_read_second_half_lse( + self, + lse: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + second_half_lse_seqlen: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_second_half_lse(*args, **kwargs) - - def thd_out_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_second_half_lse( + lse, cu_seqlens, lse_packed, second_half_lse_seqlen + ) + def thd_out_correction( + self, + out: torch.Tensor, + out_per_step: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + only_second_half: bool, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_out_correction(*args, **kwargs) - - def thd_grad_correction(self, *args, **kwargs) -> Any: + return tex.thd_out_correction( + out, out_per_step, lse, lse_per_step, + cu_seqlens, only_second_half, lse_packed + ) + def thd_grad_correction( + self, + grad: torch.Tensor, + grad_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + first_half: str, + second_half: str, + ) -> None: tex = self._get_tex() - return tex.thd_grad_correction(*args, **kwargs) - - def thd_get_partitioned_indices(self, *args, **kwargs) -> Any: + return tex.thd_grad_correction( + grad, grad_per_step, cu_seqlens, + first_half, second_half + ) + def thd_get_partitioned_indices( + self, + cu_seqlens: torch.Tensor, + total_tokens: int, + world_size: int, + rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_get_partitioned_indices(*args, **kwargs) + return tex.thd_get_partitioned_indices( + cu_seqlens, total_tokens, world_size, rank + ) - def init_nvshmem_backend(self, *args, **kwargs) -> None: + # nvshmem functions + def init_nvshmem_backend( + self, + process_group: Any, + ) -> None: tex = self._get_tex() - tex.init_nvshmem_backend(*args, **kwargs) - - def create_nvshmem_tensor(self, *args, **kwargs) -> torch.Tensor: + return tex.init_nvshmem_backend(process_group) + def create_nvshmem_tensor( + self, + shape: List[int], + dtype: torch.dtype, + ) -> torch.Tensor: tex = self._get_tex() - return tex.create_nvshmem_tensor(*args, **kwargs) - - def nvshmem_send_on_current_stream(self, *args, **kwargs) -> None: + return tex.create_nvshmem_tensor(shape, dtype) + def nvshmem_send_on_current_stream( + self, + src: torch.Tensor, + dst: torch.Tensor, + peer: int, + signal: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.nvshmem_send_on_current_stream(*args, **kwargs) - - def nvshmem_wait_on_current_stream(self, *args, **kwargs) -> None: + return tex.nvshmem_send_on_current_stream(src, dst, peer, signal) + def nvshmem_wait_on_current_stream( + self, + signal: torch.Tensor, + wait_kind: str, + ) -> None: tex = self._get_tex() - tex.nvshmem_wait_on_current_stream(*args, **kwargs) - + return tex.nvshmem_wait_on_current_stream(signal, wait_kind) def nvshmem_finalize(self) -> None: tex = self._get_tex() - tex.nvshmem_finalize() + return tex.nvshmem_finalize() + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -965,98 +1200,194 @@ def multi_tensor_scale( scale: float, ) -> None: tex = self._get_tex() - tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) - + return tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.multi_tensor_l2norm(chunk_size, noop_flag, tensor_lists, per_tensor) - def multi_tensor_unscale_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - scale: torch.Tensor, - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, scale, per_tensor) - + return tex.multi_tensor_unscale_l2norm( + chunk_size, noop_flag, tensor_lists, inv_scale, per_tensor + ) def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): - tex = self._get_tex() - if chunk_size is None: - return tex.multi_tensor_adam - tex.multi_tensor_adam( - chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, - eps, step, mode, bias_correction, weight_decay + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: + tex = self._get_tex() + return tex.multi_tensor_adam( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay ) - - def multi_tensor_adam_param_remainder(self, *args, **kwargs) -> None: + def multi_tensor_adam_param_remainder( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_param_remainder(*args, **kwargs) - - def multi_tensor_adam_fp8(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_param_remainder( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay + ) + def multi_tensor_adam_fp8( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + fp8_dtype: DType, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_fp8(*args, **kwargs) - - def multi_tensor_adam_capturable(self, *args, **kwargs) -> None: + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.multi_tensor_adam_fp8( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + fp8_dtype + ) + def multi_tensor_adam_capturable( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable(*args, **kwargs) - - def multi_tensor_adam_capturable_master(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_adam_capturable_master( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable_master(*args, **kwargs) - - def multi_tensor_sgd(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable_master( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_sgd( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_sgd(*args, **kwargs) - - def multi_tensor_compute_scale_and_scale_inv(self, *args, **kwargs) -> None: + return tex.multi_tensor_sgd( + chunk_size, noop_flag, tensor_lists, + wd, momentum, dampening, + lr, nesterov, first_run, + wd_after_momentum, scale + ) + def multi_tensor_compute_scale_and_scale_inv( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + max_fp8: float, + force_pow_2_scales: bool, + epsilon: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_compute_scale_and_scale_inv(*args, **kwargs) + return tex.multi_tensor_compute_scale_and_scale_inv( + chunk_size, noop_flag, tensor_lists, + max_fp8, force_pow_2_scales, epsilon + ) + # Comm+GEMM Overlap def bulk_overlap_ag_with_external_gemm( self, - allgather_communicator: Any, + allgather_communicator: CommOverlap, send_stream: Any, recv_stream: Any, ) -> Any: tex = self._get_tex() return tex.bulk_overlap_ag_with_external_gemm(allgather_communicator, send_stream, recv_stream) +############## class func ################################# + def get_flash_attention_class(self): + raise NotImplementedError("get_flash_attention_class - not implemented in iluvatar backend") def create_fp8_tensor_meta(self) -> FP8TensorMeta: tex = self._get_tex() return tex.FP8TensorMeta() - def create_comm_overlap_helper( self, world_group: Optional[Any] = None, intra_node_group: Optional[Any] = None, - ) -> Any: + ) -> "CommOverlapHelper": tex = self._get_tex() - if world_group is None: - return tex.CommOverlapHelper() return tex.CommOverlapHelper(world_group, intra_node_group) - def create_comm_overlap( self, buffer_shape: List[int], @@ -1072,7 +1403,7 @@ def create_comm_overlap( set_sm_margin: bool = True, atomic_gemm: bool = False, rs_overlap_first_gemm: bool = False, - ) -> Any: + ) -> "CommOverlap": tex = self._get_tex() return tex.CommOverlap( buffer_shape, buffer_dtype, helper, tp_size, @@ -1080,7 +1411,6 @@ def create_comm_overlap( gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm ) - def create_comm_overlap_p2p( self, buffer_shape: List[int], @@ -1097,13 +1427,10 @@ def create_comm_overlap_p2p( atomic_gemm: bool = False, use_ce: bool = True, aggregate: bool = False, - ) -> Any: + ) -> "CommOverlapP2P": tex = self._get_tex() return tex.CommOverlapP2P( buffer_shape, buffer_dtype, helper, tp_size, comm_type, num_max_streams, comm_cga_size, gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, use_ce, aggregate ) - - - diff --git a/transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py b/transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py index 6066a53892..9d9bb164fa 100644 --- a/transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py +++ b/transformer_engine/plugin/core/backends/vendor/kunlunxin/kunlunxin.py @@ -5,10 +5,8 @@ import os import subprocess from typing import Any, Dict, List, Optional, Tuple, Union - import torch - -from transformer_engine.plugin.core.ops import TEFLBackendBase, FP8TensorMeta, NVTE_Fused_Attn_Backend +from ....ops import * _kunlunxin_available = False diff --git a/transformer_engine/plugin/core/backends/vendor/metax/metax.py b/transformer_engine/plugin/core/backends/vendor/metax/metax.py index 8efbbc9490..6b33369c75 100644 --- a/transformer_engine/plugin/core/backends/vendor/metax/metax.py +++ b/transformer_engine/plugin/core/backends/vendor/metax/metax.py @@ -14,7 +14,7 @@ import torch -from ....ops import TEFLBackendBase, FP8TensorMeta +from ....ops import * def _load_metax_libs(): @@ -74,67 +74,6 @@ def _get_tex(): import transformer_engine_torch_metax return transformer_engine_torch_metax -def _torch_dtype_to_te_dtype(torch_dtype, tex_module): - if torch_dtype is None: - return None - - NativeDType = tex_module.DType - if type(torch_dtype).__name__ == 'DType' and type(torch_dtype).__module__ == 'transformer_engine_torch_metax': - return torch_dtype - - if hasattr(torch_dtype, 'name') and hasattr(torch_dtype, 'value'): - from transformer_engine.plugin.core.ops import DType as PyDType - if isinstance(torch_dtype, PyDType): - dtype_name = torch_dtype.name - if hasattr(NativeDType, dtype_name): - return getattr(NativeDType, dtype_name) - - dtype_map = { - torch.float32: NativeDType.kFloat32, - torch.float16: NativeDType.kFloat16, - torch.bfloat16: NativeDType.kBFloat16, - torch.int32: NativeDType.kInt32, - torch.uint8: NativeDType.kByte, - } - - if hasattr(torch, 'float8_e4m3fn'): - dtype_map[torch.float8_e4m3fn] = NativeDType.kFloat8E4M3 - if hasattr(torch, 'float8_e5m2'): - dtype_map[torch.float8_e5m2] = NativeDType.kFloat8E5M2 - - return dtype_map.get(torch_dtype, torch_dtype) - -def _convert_dtype_params(func): - - @functools.wraps(func) - def wrapper(self, *args, **kwargs): - dtype_params = ['otype', 'output_dtype', 'bias_type'] - - from transformer_engine.plugin.core.ops import DType as PyDType - - def needs_conversion(val): - return isinstance(val, torch.dtype) or isinstance(val, PyDType) - - for param_name in dtype_params: - if param_name in kwargs: - value = kwargs[param_name] - if needs_conversion(value): - converted = self._to_te_dtype(value) - kwargs[param_name] = converted - - sig = inspect.signature(func) - param_names = list(sig.parameters.keys())[1:] - - args_list = list(args) - for i, (param_name, arg_value) in enumerate(zip(param_names, args_list)): - if param_name in dtype_params and needs_conversion(arg_value): - converted = self._to_te_dtype(arg_value) - args_list[i] = converted - - return func(self, *args_list, **kwargs) - - return wrapper - class MetaxBackend(TEFLBackendBase): @staticmethod def check_available() -> bool: @@ -148,16 +87,9 @@ def _get_tex(self): self._tex = _get_tex() return self._tex - def _to_te_dtype(self, torch_dtype): - return _torch_dtype_to_te_dtype(torch_dtype, self._get_tex()) - def is_available(self) -> bool: return _check_metax_available() - def get_flash_attention_class(self): - from .flash_attention import FlashAttentionMETAX - return FlashAttentionMETAX - def get_attention_backend(self, attention_params=None): # Import the metax get_attention_backend function try: @@ -175,6 +107,7 @@ def get_attention_backend(self, attention_params=None): f"Attention_params: {self.attention_params}" ) +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def quantize( self, tensor: torch.Tensor, @@ -185,35 +118,34 @@ def quantize( tex = self._get_tex() return tex.quantize(tensor, quantizer, output, noop) - @_convert_dtype_params def dequantize( self, - input: torch.Tensor, - otype: torch.dtype, - ) -> torch.Tensor: + input: Any, + otype: DType, + ) -> Any: tex = self._get_tex() + otype = tex.DType(int(otype)) if otype is not None else None return tex.dequantize(input, otype) def bgrad_quantize( self, input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: tex = self._get_tex() return tex.bgrad_quantize(input, quantizer) - @_convert_dtype_params def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -222,61 +154,53 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: + ) -> List[Any]: tex = self._get_tex() - - if bias_type is None: - bias_type = self._to_te_dtype(torch.bfloat16) - + + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + comm_type = tex.CommOverlapType(int(comm_type)) if comm_type is not None else None + output_dtype = tex.DType(int(output_dtype)) if output_dtype is not None else None return tex.generic_gemm( A, transA, B, transB, D, quantizer, output_dtype, bias, bias_type, gelu, gelu_in, grad, workspace, workspace_size, accumulate, use_split_accumulator, comm_overlap, comm_type, extra_output, bulk_overlap, alpha, beta ) - - def te_general_grouped_gemm(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.te_general_grouped_gemm(*args, **kwargs) - + # GELU and variants # def gelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.gelu(input, quantizer) - def geglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.geglu(input, quantizer) def qgelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgelu(input, quantizer) - def qgeglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.qgeglu(input, quantizer) + # ReLU and variants # def relu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.relu(input, quantizer) - def reglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.reglu(input, quantizer) def srelu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.srelu(input, quantizer) - def sreglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.sreglu(input, quantizer) - + # SwiGLU and variants # def silu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.silu(input, quantizer) - def swiglu(self, input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.swiglu(input, quantizer) @@ -289,42 +213,39 @@ def clamped_swiglu( ) -> Any: tex = self._get_tex() return tex.clamped_swiglu(input, quantizer, limit, alpha) - + # Backward of GELU and variants # def dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgelu(grad, fwd_input, quantizer) def dgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dgeglu(grad, fwd_input, quantizer) - def dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgelu(grad, fwd_input, quantizer) def dqgeglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dqgeglu(grad, fwd_input, quantizer) - + # Backward of ReLU and variants # def drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.drelu(grad, fwd_input, quantizer) def dreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dreglu(grad, fwd_input, quantizer) - def dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsrelu(grad, fwd_input, quantizer) def dsreglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsreglu(grad, fwd_input, quantizer) - + # Backward of SiLU and variants # def dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dsilu(grad, fwd_input, quantizer) def dswiglu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Any: tex = self._get_tex() return tex.dswiglu(grad, fwd_input, quantizer) - def clamped_dswiglu( self, grad: torch.Tensor, @@ -335,131 +256,207 @@ def clamped_dswiglu( ) -> Any: tex = self._get_tex() return tex.clamped_dswiglu(grad, fwd_input, quantizer, limit, alpha) - - def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + # DBias + DAct fusions # + def dbias_dgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dgelu(grad, fwd_input, quantizer) - - def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsilu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsilu(grad, fwd_input, quantizer) - - def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_drelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_drelu(grad, fwd_input, quantizer) - - def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dqgelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dqgelu(grad, fwd_input, quantizer) - - def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> Tuple[torch.Tensor, Any]: + def dbias_dsrelu(self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any) -> List[Any]: tex = self._get_tex() return tex.dbias_dsrelu(grad, fwd_input, quantizer) - - @_convert_dtype_params + # Permutation functions + def moe_permute_fwd( + self, + input: torch.Tensor, + dtype: DType, + indices: torch.Tensor, + num_out_tokens: int, + workspace: List[torch.Tensor], + max_expanded_token_num: int, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_fwd(input, dtype,indices,num_out_tokens,workspace,max_expanded_token_num) + def moe_permute_bwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_permute_bwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_fwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_fwd(input,dtype,row_id_map,prob,num_tokens,topK) + def moe_unpermute_bwd( + self, + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + tex = self._get_tex() + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.moe_unpermute_bwd(input_bwd,input_fwd,dtype,row_id_map,prob) + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_forward(input, scale) + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_forward(input, mask, scale_factor) + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_masked_softmax_backward(output_grad_, softmax_results_, scale_factor) + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_forward(input, scale_factor) + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_upper_triang_masked_softmax_backward( + output_grads_, softmax_results_, scale_factor + ) + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_forward(input, scale_factor) + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + tex = self._get_tex() + return tex.scaled_aligned_causal_masked_softmax_backward( + output_grad_, softmax_results_, scale_factor + ) + # Other granular functions def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, mu, rsigma = tex.layernorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.layernorm_fwd( input, weight, bias, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - return y, mu, rsigma - def layernorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, mu: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dgamma, dbeta = tex.layernorm_bwd(dy, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dgamma, dbeta - - @_convert_dtype_params + return tex.layernorm_bwd( + dz, x, mu, rsigma, gamma, sm_margin, zero_centered_gamma + ) def rmsnorm_fwd( self, - input: torch.Tensor, - weight: torch.Tensor, + input: Any, + weight: Any, eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: tex = self._get_tex() - - orig_shape = input.shape - if input.ndim > 2: - input = input.view(-1, input.shape[-1]) - - y, y_quant, rsigma = tex.rmsnorm_fwd( + otype = tex.DType(int(otype)) if otype is not None else None + return tex.rmsnorm_fwd( input, weight, eps, ln_out, quantizer, otype, sm_margin, zero_centered_gamma ) - - if len(orig_shape) > 2: - y = y.view(*orig_shape) - if y_quant is not None: - y_quant = y_quant.view(*orig_shape) - return y, y_quant, rsigma - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - - orig_shape = dy.shape - if dy.ndim > 2: - dy = dy.view(-1, dy.shape[-1]) - x = x.view(-1, x.shape[-1]) - - dx, dw = tex.rmsnorm_bwd(dy, x, rsigma, gamma, sm_margin, zero_centered_gamma) - - if len(orig_shape) > 2: - dx = dx.view(*orig_shape) - return dx, dw - - def rmsnorm_bwd_add(self, *args, **kwargs) -> Any: + return tex.rmsnorm_bwd(dz, x, rsigma, gamma, sm_margin, zero_centered_gamma) + def rmsnorm_bwd_add( + self, + dz: torch.Tensor, + x: torch.Tensor, + add: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.rmsnorm_bwd_add(*args, **kwargs) + return tex.rmsnorm_bwd_add(dz, x, add, rsigma, gamma, sm_margin, zero_centered_gamma) def multi_tensor_quantize( self, @@ -468,7 +465,6 @@ def multi_tensor_quantize( ) -> List[Any]: tex = self._get_tex() return tex.multi_tensor_quantize(tensor_list, quantizer_list) - def split_quantize( self, tensor: torch.Tensor, @@ -477,246 +473,457 @@ def split_quantize( ) -> List[Any]: tex = self._get_tex() return tex.split_quantize(tensor, split_sections, quantizer_list) - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_fwd(*args, **kwargs) - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_permute_bwd(*args, **kwargs) - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_fwd(*args, **kwargs) - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.moe_unpermute_bwd(*args, **kwargs) - - def scaled_softmax_forward(self, input: torch.Tensor, scale: float) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_forward(input, scale) - - def scaled_softmax_backward( + def te_general_grouped_gemm( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: - tex = self._get_tex() - return tex.scaled_softmax_backward(output_grad, softmax_output, scale) - - def scaled_masked_softmax_forward( + A: List[Any], + transa: bool, + B: List[Any], + transb: bool, + D: Optional[List[torch.Tensor]], + D_type: DType, + m_splits: List[int], + bias: List[torch.Tensor], + bias_type: DType, + single_output: bool, + pre_gelu_out: List[torch.Tensor], + grad: bool, + workspace: List[torch.Tensor], + workspaceSizes: int, + accumulate: bool, + use_split_accumulator: bool, + math_sm_count: int, + ) -> Optional[List[torch.Tensor]]: + tex = self._get_tex() + D_type = tex.DType(int(D_type)) if D_type is not None else None + bias_type = tex.DType(int(bias_type)) if bias_type is not None else None + return tex.te_general_grouped_gemm( + A, transa, B, transb, D, D_type, m_splits, bias, bias_type, + single_output, pre_gelu_out, grad, workspace, workspaceSizes, + accumulate, use_split_accumulator, math_sm_count + ) + def fp8_transpose( self, input: torch.Tensor, - mask: torch.Tensor, - scale: float, + dtype: DType, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_forward(input, mask, scale) - - def scaled_masked_softmax_backward( + dtype = tex.DType(int(dtype)) if dtype is not None else None + return tex.fp8_transpose(input, dtype, out) + def swap_first_dims( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + tensor: torch.Tensor, + out: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() - return tex.scaled_masked_softmax_backward(output_grad, softmax_output, scale) + return tex.swap_first_dims(tensor, out) + def get_fused_attn_backend( + self, + is_training: bool, + q_dtype: DType, + kv_dtype: DType, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + p_dropout: float, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + window_size_left: int, + window_size_right: int, + return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: + tex = self._get_tex() + + q_dtype = tex.DType(int(q_dtype)) if q_dtype is not None else None + kv_dtype = tex.DType(int(kv_dtype)) if kv_dtype is not None else None + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + result = tex.get_fused_attn_backend( + is_training, q_dtype, kv_dtype, qkv_layout, bias_type, + attn_mask_type, softmax_type, p_dropout, num_attn_heads, + num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, + head_dim_v, window_size_left, window_size_right, return_max_logit + ) + return NVTE_Fused_Attn_Backend(result) - def scaled_upper_triang_masked_softmax_forward( + def compute_amax( self, input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax: torch.Tensor, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_forward(input, scale) - - def scaled_upper_triang_masked_softmax_backward( + return tex.compute_amax(input, amax) + def fused_amax_and_scale_update_after_reduction( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax_reduction_buffer: torch.Tensor, + amax_histories: List[torch.Tensor], + scales: List[torch.Tensor], + amax_compute_algo: str, + fp8_dtype: DType, + margin: float, + ) -> None: tex = self._get_tex() - return tex.scaled_upper_triang_masked_softmax_backward(output_grad, softmax_output, scale) - - def scaled_aligned_causal_masked_softmax_forward( + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.fused_amax_and_scale_update_after_reduction( + amax_reduction_buffer, amax_histories, scales, + amax_compute_algo, fp8_dtype, margin + ) + def fp8_block_scaling_compute_partial_amax( self, - input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_forward(input, scale) - - def scaled_aligned_causal_masked_softmax_backward( + return tex.fp8_block_scaling_compute_partial_amax( + tensor, amax, h, w, start_offset, block_len + ) + def fp8_block_scaling_partial_cast( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + inp: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + out_dtype: DType, + ) -> None: tex = self._get_tex() - return tex.scaled_aligned_causal_masked_softmax_backward(output_grad, softmax_output, scale) - - def get_fused_attn_backend(self, *args, **kwargs) -> int: + out_dtype = tex.DType(int(out_dtype)) if out_dtype is not None else None + return tex.fp8_block_scaling_partial_cast( + inp, out, scale, h, w, start_offset, block_len, out_dtype + ) + def fused_multi_row_padding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + padded_input_row_list: List[int], + ) -> None: tex = self._get_tex() - - args_list = list(args) - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - - if type(py_enum).__module__ == 'transformer_engine_torch_metax': - return py_enum - - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - - if hasattr(py_enum, 'value'): - enum_value = int(py_enum.value) - for member_name in dir(native_enum_class): - if not member_name.startswith('_'): - try: - member = getattr(native_enum_class, member_name) - if hasattr(member, 'value') and int(member.value) == enum_value: - return member - except: - pass - - if hasattr(py_enum, 'value'): - return int(py_enum.value) - - return py_enum - - if len(args) > 1: - args_list[1] = self._to_te_dtype(args[1]) - if len(args) > 2: - args_list[2] = self._to_te_dtype(args[2]) - if len(args) > 3: - args_list[3] = convert_enum(args[3], tex.NVTE_QKV_Layout) - if len(args) > 4: - args_list[4] = convert_enum(args[4], tex.NVTE_Bias_Type) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_Mask_Type) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Softmax_Type) - - return tex.get_fused_attn_backend(*args_list, **kwargs) - - def fused_attn_fwd(self, *args, **kwargs) -> Any: + return tex.fused_multi_row_padding( + input, output, input_row_list, padded_input_row_list + ) + def fused_multi_row_unpadding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + unpadded_input_row_list: List[int], + ) -> None: tex = self._get_tex() + return tex.fused_multi_row_unpadding( + input, output, input_row_list, unpadded_input_row_list + ) - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_metax': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_QKV_Layout) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Bias_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Mask_Type) - if len(args) > 9: - args_list[9] = convert_enum(args[9], tex.NVTE_Softmax_Type) - - return tex.fused_attn_fwd(*args_list, **kwargs) - - def fused_attn_bwd(self, *args, **kwargs) -> Any: + # attention kernels + def fa_prepare_fwd( + self, + qkvi: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - - def convert_enum(py_enum, native_enum_class): - if py_enum is None: - return None - if type(py_enum).__module__ == 'transformer_engine_torch_metax': - return py_enum - if hasattr(py_enum, 'name'): - enum_name = py_enum.name - if hasattr(native_enum_class, enum_name): - return getattr(native_enum_class, enum_name) - return py_enum - - args_list = list(args) - if len(args) > 5: - args_list[5] = convert_enum(args[5], tex.NVTE_QKV_Layout) - if len(args) > 6: - args_list[6] = convert_enum(args[6], tex.NVTE_Bias_Type) - if len(args) > 7: - args_list[7] = convert_enum(args[7], tex.NVTE_Mask_Type) - if len(args) > 8: - args_list[8] = convert_enum(args[8], tex.NVTE_Softmax_Type) - if len(args) > 19: - args_list[19] = self._to_te_dtype(args[19]) - - if 'dqkv_dtype' in kwargs: - kwargs['dqkv_dtype'] = self._to_te_dtype(kwargs['dqkv_dtype']) - - return tex.fused_attn_bwd(*args_list, **kwargs) - - def fa_prepare_fwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_fwd(qkvi) + def fa_prepare_bwd( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fa_prepare_fwd(*args, **kwargs) - - def fa_prepare_bwd(self, *args, **kwargs) -> Any: + return tex.fa_prepare_bwd(q, k, v) + def fused_attn_fwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + is_training: bool, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + fake_dtype: torch.dtype, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + s_quantizer: Any, + o_quantizer: Any, + Bias: Optional[torch.Tensor], + SoftmaxOffset: Optional[torch.Tensor], + rng_gen: Optional[torch.Generator], + rng_elts_per_thread: int, + return_max_logit: bool, + ) -> List[Any]: tex = self._get_tex() - return tex.fa_prepare_bwd(*args, **kwargs) - def copy_to_kv_cache(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + + return tex.fused_attn_fwd( + max_seqlen_q, + max_seqlen_kv, + is_training, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + fake_dtype, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + page_table_k, + page_table_v, + s_quantizer, + o_quantizer, + Bias, + SoftmaxOffset, + rng_gen, + rng_elts_per_thread, + return_max_logit + ) + def fused_attn_bwd( + self, + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + O: Any, + dO: Any, + fake_dtype: torch.dtype, + dqkv_type: DType, + Aux_CTX_Tensors: List[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + s_quantizer: Any, + dp_quantizer: Any, + dqkv_quantizer: Any, + ) -> List[Any]: tex = self._get_tex() - return tex.copy_to_kv_cache(*args, **kwargs) - def convert_thd_to_bshd(self, *args, **kwargs) -> Any: + qkv_layout = tex.NVTE_QKV_Layout(int(qkv_layout)) if qkv_layout is not None else None + bias_type = tex.NVTE_Bias_Type(int(bias_type)) if bias_type is not None else None + attn_mask_type = tex.NVTE_Mask_Type(int(attn_mask_type)) if attn_mask_type is not None else None + softmax_type = tex.NVTE_Softmax_Type(int(softmax_type)) if softmax_type is not None else None + dqkv_type = tex.DType(int(dqkv_type)) if dqkv_type is not None else None + + return tex.fused_attn_bwd( + max_seqlen_q, + max_seqlen_kv, + attn_scale, + p_dropout, + set_zero, + qkv_layout, + bias_type, + attn_mask_type, + softmax_type, + window_size, + deterministic, + cu_seqlens_q, + cu_seqlens_kv, + Q, + K, + V, + O, + dO, + fake_dtype, + dqkv_type, + Aux_CTX_Tensors, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, + s_quantizer, + dp_quantizer, + dqkv_quantizer + ) + def copy_to_kv_cache( + self, + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: NVTE_QKV_Format, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, + ) -> None: tex = self._get_tex() - return tex.convert_thd_to_bshd(*args, **kwargs) - - def convert_bshd_to_thd(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.copy_to_kv_cache( + new_k, + new_v, + k_cache, + v_cache, + page_table, + cu_new_lens, + cu_cached_lens, + qkv_format, + b, + max_ctx_len, + max_seq_len, + max_pages_per_seq, + is_non_paged + ) + def convert_thd_to_bshd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + b: int, + max_seq_len: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.convert_bshd_to_thd(*args, **kwargs) - - def fused_rope_forward(self, *args, **kwargs) -> Any: + return tex.convert_thd_to_bshd(tensor, cu_seqlens, b, max_seq_len) + def convert_bshd_to_thd( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + t: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_rope_forward(*args, **kwargs) + return tex.convert_bshd_to_thd(tensor, cu_seqlens, t) - def fused_rope_backward(self, *args, **kwargs) -> Any: + # 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: tex = self._get_tex() - return tex.fused_rope_backward(*args, **kwargs) - - def fused_qkv_rope_forward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_forward( + 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, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_qkv_rope_forward(*args, **kwargs) - - def fused_qkv_rope_backward(self, *args, **kwargs) -> Any: + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_rope_backward( + output_grads, freqs, 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]: + tex = self._get_tex() + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.fused_qkv_rope_forward( + 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: tex = self._get_tex() - return tex.fused_qkv_rope_backward(*args, **kwargs) + qkv_format = tex.NVTE_QKV_Format(int(qkv_format)) if qkv_format is not None else None + return tex.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 + ) + # fused router def fused_topk_with_score_function_fwd( self, logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, - score_function: Any, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, expert_bias: Optional[torch.Tensor], - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_topk_with_score_function_fwd( - logits, topk, use_pre_softmax, num_groups, group_topk, - scaling_factor, score_function, expert_bias + logits, + topk, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + score_function, + expert_bias, ) - def fused_topk_with_score_function_bwd( self, num_tokens: int, @@ -726,24 +933,33 @@ def fused_topk_with_score_function_bwd( grad_probs: torch.Tensor, topk: int, use_pre_softmax: bool, - scaling_factor: float, - score_function: Any, - ) -> Any: + scaling_factor: Optional[float], + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_topk_with_score_function_bwd( - num_tokens, num_experts, routing_map, intermediate_output, - grad_probs, topk, use_pre_softmax, scaling_factor, score_function + num_tokens, + num_experts, + routing_map, + intermediate_output, + grad_probs, + topk, + use_pre_softmax, + scaling_factor, + score_function, ) - def fused_score_for_moe_aux_loss_fwd( self, logits: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.fused_score_for_moe_aux_loss_fwd(logits, topk, score_function) - + return tex.fused_score_for_moe_aux_loss_fwd( + logits, + topk, + score_function, + ) def fused_score_for_moe_aux_loss_bwd( self, num_tokens: int, @@ -751,13 +967,17 @@ def fused_score_for_moe_aux_loss_bwd( intermediate_output: torch.Tensor, grad_scores: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> torch.Tensor: tex = self._get_tex() return tex.fused_score_for_moe_aux_loss_bwd( - num_tokens, num_experts, intermediate_output, grad_scores, topk, score_function + num_tokens, + num_experts, + intermediate_output, + grad_scores, + topk, + score_function, ) - def fused_moe_aux_loss_fwd( self, probs: torch.Tensor, @@ -768,13 +988,18 @@ def fused_moe_aux_loss_fwd( num_cols: int, topk: int, coeff: float, - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.fused_moe_aux_loss_fwd( - probs, tokens_per_expert, total_num_tokens, num_experts, - num_rows, num_cols, topk, coeff + probs, + tokens_per_expert, + total_num_tokens, + num_experts, + num_rows, + num_cols, + topk, + coeff, ) - def fused_moe_aux_loss_bwd( self, Const_buf: torch.Tensor, @@ -782,152 +1007,146 @@ def fused_moe_aux_loss_bwd( num_rows: int, num_cols: int, grad_aux_loss: torch.Tensor, - ) -> Any: + ) -> torch.Tensor: tex = self._get_tex() - return tex.fused_moe_aux_loss_bwd( - Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss - ) + return tex.fused_moe_aux_loss_bwd(Const_buf, tokens_per_expert, num_rows, num_cols, grad_aux_loss) + # Dropout def dropout_fwd( self, input: torch.Tensor, dropout_probability: float, - out: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.dropout_fwd(input, dropout_probability, out) - def dropout_bwd( self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, - grad_input: Optional[torch.Tensor] = None, + grad_input: Optional[torch.Tensor], ) -> torch.Tensor: tex = self._get_tex() return tex.dropout_bwd(grad_output, mask, dropout_probability, grad_input) - def fp8_transpose( - self, - input: torch.Tensor, - dtype: Any, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.fp8_transpose(input, dtype, out=out) - - def swap_first_dims( - self, - tensor: torch.Tensor, - *, - out: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.swap_first_dims(tensor, out=out) - - def compute_amax( - self, - input: torch.Tensor, - amax: torch.Tensor, - ) -> None: - tex = self._get_tex() - tex.compute_amax(input, amax) - - def fused_amax_and_scale_update_after_reduction(self, *args, **kwargs) -> None: - tex = self._get_tex() - tex.fused_amax_and_scale_update_after_reduction(*args, **kwargs) - - def fp8_block_scaling_compute_partial_amax( - self, - tensor: torch.Tensor, - amax: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_compute_partial_amax(tensor, amax, h, w, start_offset, block_len) - - def fp8_block_scaling_partial_cast( - self, - inp: torch.Tensor, - out: torch.Tensor, - scale: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - out_dtype: Any, - ) -> None: - tex = self._get_tex() - tex.fp8_block_scaling_partial_cast(inp, out, scale, h, w, start_offset, block_len, out_dtype) - - def fused_multi_row_padding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_padding(*args, **kwargs) - - def fused_multi_row_unpadding(self, *args, **kwargs) -> Any: - tex = self._get_tex() - return tex.fused_multi_row_unpadding(*args, **kwargs) - + # Misc def get_cublasLt_version(self) -> int: tex = self._get_tex() return tex.get_cublasLt_version() - def get_cudnn_version(self) -> int: tex = self._get_tex() return tex.get_cudnn_version() - def get_num_cublas_streams(self) -> int: tex = self._get_tex() return tex.get_num_cublas_streams() - def thd_read_half_tensor(self, *args, **kwargs) -> Any: + # Support THD format for Context Parallel + def thd_read_half_tensor( + self, + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + half_idx: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_half_tensor(*args, **kwargs) - - def thd_second_half_lse_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_half_tensor(tensor, cu_seqlens, half_idx) + def thd_second_half_lse_correction( + self, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_second_half_lse_correction(*args, **kwargs) - - def thd_read_second_half_lse(self, *args, **kwargs) -> Any: + return tex.thd_second_half_lse_correction( + lse, lse_per_step, cu_seqlens, lse_packed + ) + def thd_read_second_half_lse( + self, + lse: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + second_half_lse_seqlen: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_read_second_half_lse(*args, **kwargs) - - def thd_out_correction(self, *args, **kwargs) -> Any: + return tex.thd_read_second_half_lse( + lse, cu_seqlens, lse_packed, second_half_lse_seqlen + ) + def thd_out_correction( + self, + out: torch.Tensor, + out_per_step: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + only_second_half: bool, + lse_packed: bool, + ) -> None: tex = self._get_tex() - return tex.thd_out_correction(*args, **kwargs) - - def thd_grad_correction(self, *args, **kwargs) -> Any: + return tex.thd_out_correction( + out, out_per_step, lse, lse_per_step, + cu_seqlens, only_second_half, lse_packed + ) + def thd_grad_correction( + self, + grad: torch.Tensor, + grad_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + first_half: str, + second_half: str, + ) -> None: tex = self._get_tex() - return tex.thd_grad_correction(*args, **kwargs) - - def thd_get_partitioned_indices(self, *args, **kwargs) -> Any: + return tex.thd_grad_correction( + grad, grad_per_step, cu_seqlens, + first_half, second_half + ) + def thd_get_partitioned_indices( + self, + cu_seqlens: torch.Tensor, + total_tokens: int, + world_size: int, + rank: int, + ) -> torch.Tensor: tex = self._get_tex() - return tex.thd_get_partitioned_indices(*args, **kwargs) + return tex.thd_get_partitioned_indices( + cu_seqlens, total_tokens, world_size, rank + ) - def init_nvshmem_backend(self, *args, **kwargs) -> None: + # nvshmem functions + def init_nvshmem_backend( + self, + process_group: Any, + ) -> None: tex = self._get_tex() - tex.init_nvshmem_backend(*args, **kwargs) - - def create_nvshmem_tensor(self, *args, **kwargs) -> torch.Tensor: + return tex.init_nvshmem_backend(process_group) + def create_nvshmem_tensor( + self, + shape: List[int], + dtype: torch.dtype, + ) -> torch.Tensor: tex = self._get_tex() - return tex.create_nvshmem_tensor(*args, **kwargs) - - def nvshmem_send_on_current_stream(self, *args, **kwargs) -> None: + return tex.create_nvshmem_tensor(shape, dtype) + def nvshmem_send_on_current_stream( + self, + src: torch.Tensor, + dst: torch.Tensor, + peer: int, + signal: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.nvshmem_send_on_current_stream(*args, **kwargs) - - def nvshmem_wait_on_current_stream(self, *args, **kwargs) -> None: + return tex.nvshmem_send_on_current_stream(src, dst, peer, signal) + def nvshmem_wait_on_current_stream( + self, + signal: torch.Tensor, + wait_kind: str, + ) -> None: tex = self._get_tex() - tex.nvshmem_wait_on_current_stream(*args, **kwargs) - + return tex.nvshmem_wait_on_current_stream(signal, wait_kind) def nvshmem_finalize(self) -> None: tex = self._get_tex() - tex.nvshmem_finalize() + return tex.nvshmem_finalize() + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -936,98 +1155,195 @@ def multi_tensor_scale( scale: float, ) -> None: tex = self._get_tex() - tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) - + return tex.multi_tensor_scale(chunk_size, noop_flag, tensor_lists, scale) def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() return tex.multi_tensor_l2norm(chunk_size, noop_flag, tensor_lists, per_tensor) - def multi_tensor_unscale_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - scale: torch.Tensor, - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: tex = self._get_tex() - return tex.multi_tensor_unscale_l2norm(chunk_size, noop_flag, tensor_lists, scale, per_tensor) - + return tex.multi_tensor_unscale_l2norm( + chunk_size, noop_flag, tensor_lists, inv_scale, per_tensor + ) def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): - tex = self._get_tex() - if chunk_size is None: - return tex.multi_tensor_adam - tex.multi_tensor_adam( - chunk_size, noop_flag, tensor_lists, lr, beta1, beta2, - eps, step, mode, bias_correction, weight_decay + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: + tex = self._get_tex() + return tex.multi_tensor_adam( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay ) - - def multi_tensor_adam_param_remainder(self, *args, **kwargs) -> None: + def multi_tensor_adam_param_remainder( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_param_remainder(*args, **kwargs) - - def multi_tensor_adam_fp8(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_param_remainder( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay + ) + def multi_tensor_adam_fp8( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + fp8_dtype: DType, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_fp8(*args, **kwargs) - - def multi_tensor_adam_capturable(self, *args, **kwargs) -> None: + fp8_dtype = tex.DType(int(fp8_dtype)) if fp8_dtype is not None else None + return tex.multi_tensor_adam_fp8( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + fp8_dtype + ) + def multi_tensor_adam_capturable( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable(*args, **kwargs) - - def multi_tensor_adam_capturable_master(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_adam_capturable_master( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, + ) -> None: tex = self._get_tex() - tex.multi_tensor_adam_capturable_master(*args, **kwargs) - - def multi_tensor_sgd(self, *args, **kwargs) -> None: + return tex.multi_tensor_adam_capturable_master( + chunk_size, noop_flag, tensor_lists, + lr, beta1, beta2, epsilon, + step, mode, bias_correction, weight_decay, + inv_scale + ) + def multi_tensor_sgd( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_sgd(*args, **kwargs) - - def multi_tensor_compute_scale_and_scale_inv(self, *args, **kwargs) -> None: + return tex.multi_tensor_sgd( + chunk_size, noop_flag, tensor_lists, + wd, momentum, dampening, + lr, nesterov, first_run, + wd_after_momentum, scale + ) + def multi_tensor_compute_scale_and_scale_inv( + self, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + max_fp8: float, + force_pow_2_scales: bool, + epsilon: float, + ) -> None: tex = self._get_tex() - tex.multi_tensor_compute_scale_and_scale_inv(*args, **kwargs) + return tex.multi_tensor_compute_scale_and_scale_inv( + chunk_size, noop_flag, tensor_lists, + max_fp8, force_pow_2_scales, epsilon + ) + # Comm+GEMM Overlap def bulk_overlap_ag_with_external_gemm( self, - allgather_communicator: Any, + allgather_communicator: CommOverlap, send_stream: Any, recv_stream: Any, ) -> Any: tex = self._get_tex() return tex.bulk_overlap_ag_with_external_gemm(allgather_communicator, send_stream, recv_stream) +############## class func ################################# + def get_flash_attention_class(self): + from .flash_attention import FlashAttentionMETAX + return FlashAttentionMETAX def create_fp8_tensor_meta(self) -> FP8TensorMeta: tex = self._get_tex() return tex.FP8TensorMeta() - def create_comm_overlap_helper( self, world_group: Optional[Any] = None, intra_node_group: Optional[Any] = None, - ) -> Any: + ) -> "CommOverlapHelper": tex = self._get_tex() - if world_group is None: - return tex.CommOverlapHelper() return tex.CommOverlapHelper(world_group, intra_node_group) - def create_comm_overlap( self, buffer_shape: List[int], @@ -1043,7 +1359,7 @@ def create_comm_overlap( set_sm_margin: bool = True, atomic_gemm: bool = False, rs_overlap_first_gemm: bool = False, - ) -> Any: + ) -> "CommOverlap": tex = self._get_tex() return tex.CommOverlap( buffer_shape, buffer_dtype, helper, tp_size, @@ -1051,7 +1367,6 @@ def create_comm_overlap( gemm_priority, comm_priority, num_comm_sm, set_sm_margin, atomic_gemm, rs_overlap_first_gemm ) - def create_comm_overlap_p2p( self, buffer_shape: List[int], @@ -1068,7 +1383,7 @@ def create_comm_overlap_p2p( atomic_gemm: bool = False, use_ce: bool = True, aggregate: bool = False, - ) -> Any: + ) -> "CommOverlapP2P": tex = self._get_tex() return tex.CommOverlapP2P( buffer_shape, buffer_dtype, helper, tp_size, comm_type, diff --git a/transformer_engine/plugin/core/ops.py b/transformer_engine/plugin/core/ops.py index 988829b98c..74357394e8 100644 --- a/transformer_engine/plugin/core/ops.py +++ b/transformer_engine/plugin/core/ops.py @@ -11,6 +11,7 @@ from .logger_manager import get_logger logger = get_logger() +################### Enums ################### class DType(IntEnum): kByte = 0 kInt16 = 1 @@ -141,94 +142,260 @@ class CommOverlapAlgo(IntEnum): ATOMIC_GEMM_RS_P2P = 7 EXTERNAL_BULK_OVERLAP_AG = 8 -class FP8TensorMeta: - def __init__(self): - self.scale: Optional[torch.Tensor] = None - self.scale_inv: Optional[torch.Tensor] = None - self.amax_history: Optional[torch.Tensor] = None - -class CommGemmOverlapAlgoConfig: - def __init__(self, *args, **kwargs): - pass - -class FusedAdamCUDAKernel: - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "FusedAdamCUDAKernel requires CUDA extensions. " - "Not supported in FL mode." - ) +############ Class ################# -class FusedSGDCUDAKernel: - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "FusedSGDCUDAKernel requires CUDA extensions. " - "Not supported in FL mode." - ) +class FP8TensorMeta: + """ + FP8TensorMeta wrapper that routes to the appropriate backend implementation. + """ + def __new__(cls, *args, **kwargs): + from .manager import get_default_manager + return get_default_manager().call("create_fp8_tensor_meta", *args, **kwargs) class CommOverlapHelper: - def __init__(self, world_group=None, intra_node_group=None): - self.world_group = world_group - self.intra_node_group = intra_node_group + """ + CommOverlapHelper wrapper that routes to the appropriate backend implementation. + """ + def __new__(cls, *args, **kwargs): + from .manager import get_default_manager + return get_default_manager().call("create_comm_overlap_helper", *args, **kwargs) class CommOverlap: - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "CommOverlap should be created via backend.create_comm_overlap(). " - "Direct instantiation is not supported in FL mode." - ) + """ + CommOverlap wrapper that routes to the appropriate backend implementation. + """ + def __new__(cls, *args, **kwargs): + from .manager import get_default_manager + return get_default_manager().call("create_comm_overlap", *args, **kwargs) class CommOverlapP2P: - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "CommOverlapP2P should be created via backend.create_comm_overlap_p2p(). " - "Direct instantiation is not supported in FL mode." + """ + CommOverlapP2P wrapper that routes to the appropriate backend implementation. + """ + def __new__(cls, *args, **kwargs): + from .manager import get_default_manager + return get_default_manager().call("create_comm_overlap_p2p", *args, **kwargs) + +class FlashAttentionBase(torch.nn.Module, ABC): + def __init__( + self, + softmax_scale: float, + attention_dropout: float = 0.0, + attention_dropout_ctx: Optional[Callable] = None, + attention_type: str = "self", + layer_number: Optional[int] = None, + deterministic: bool = False, + ) -> None: + super().__init__() + + self.softmax_scale = softmax_scale + self.attention_dropout = attention_dropout + self.attention_dropout_ctx = attention_dropout_ctx or nullcontext + self.attention_type = attention_type + self.layer_number = 1 if layer_number is None else layer_number + self.deterministic = deterministic + + # For fallback support + self._manager = None + self._init_params = None + + @abstractmethod + def _forward_impl( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, + qkv_layout: str = "sbh3d", + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, + alibi_slopes: Optional[torch.Tensor] = None, + cp_group: Optional[Any] = None, + cp_global_ranks: Optional[List[int]] = None, + cp_stream: Optional[torch.cuda.Stream] = None, + cp_comm_type: str = "p2p", + fp8: bool = False, + fp8_meta: Optional[Dict[str, Any]] = None, + quantizers: Optional[Any] = None, + inference_params: Optional[Any] = None, + flash_attention_backend: Optional[Any] = None, + fp8_output: bool = False, + ) -> torch.Tensor: + """ + Actual forward implementation - subclasses must implement this. + + This method contains the backend-specific logic for flash attention. + """ + raise NotImplementedError("Subclasses must implement _forward_impl()") + + def forward( + self, + query_layer: torch.Tensor, + key_layer: torch.Tensor, + value_layer: torch.Tensor, + attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, + qkv_layout: str = "sbh3d", + cu_seqlens_q: Optional[torch.Tensor] = None, + cu_seqlens_kv: Optional[torch.Tensor] = None, + max_seqlen_q: Optional[int] = None, + max_seqlen_kv: Optional[int] = None, + attn_mask_type: str = "causal", + window_size: Optional[Tuple[int, int]] = None, + alibi_slopes: Optional[torch.Tensor] = None, + cp_group: Optional[Any] = None, + cp_global_ranks: Optional[List[int]] = None, + cp_stream: Optional[torch.cuda.Stream] = None, + cp_comm_type: str = "p2p", + fp8: bool = False, + fp8_meta: Optional[Dict[str, Any]] = None, + quantizers: Optional[Any] = None, + inference_params: Optional[Any] = None, + flash_attention_backend: Optional[Any] = None, + fp8_output: bool = False, + ) -> torch.Tensor: + """ + Forward pass with automatic fallback support and caching. + Delegates to OpManager.call_with_custom_impl for unified dispatch. + """ + if self._manager is None: + return self._forward_impl( + query_layer=query_layer, + key_layer=key_layer, + value_layer=value_layer, + attention_mask=attention_mask, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + window_size=window_size, + alibi_slopes=alibi_slopes, + cp_group=cp_group, + cp_global_ranks=cp_global_ranks, + cp_stream=cp_stream, + cp_comm_type=cp_comm_type, + fp8=fp8, + fp8_meta=fp8_meta, + quantizers=quantizers, + inference_params=inference_params, + flash_attention_backend=flash_attention_backend, + fp8_output=fp8_output, + ) + + def call_impl_fn(impl_class): + if impl_class == self.__class__: + return self._forward_impl( + query_layer=query_layer, + key_layer=key_layer, + value_layer=value_layer, + attention_mask=attention_mask, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + window_size=window_size, + alibi_slopes=alibi_slopes, + cp_group=cp_group, + cp_global_ranks=cp_global_ranks, + cp_stream=cp_stream, + cp_comm_type=cp_comm_type, + fp8=fp8, + fp8_meta=fp8_meta, + quantizers=quantizers, + inference_params=inference_params, + flash_attention_backend=flash_attention_backend, + fp8_output=fp8_output, + ) + else: + fallback_instance = impl_class(**self._init_params) + fallback_instance._manager = self._manager + fallback_instance._init_params = self._init_params + return fallback_instance._forward_impl( + query_layer=query_layer, + key_layer=key_layer, + value_layer=value_layer, + attention_mask=attention_mask, + qkv_layout=qkv_layout, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_kv, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_kv, + attn_mask_type=attn_mask_type, + window_size=window_size, + alibi_slopes=alibi_slopes, + cp_group=cp_group, + cp_global_ranks=cp_global_ranks, + cp_stream=cp_stream, + cp_comm_type=cp_comm_type, + fp8=fp8, + fp8_meta=fp8_meta, + quantizers=quantizers, + inference_params=inference_params, + flash_attention_backend=flash_attention_backend, + fp8_output=fp8_output, + ) + + return self._manager.call_with_custom_impl( + op_name="get_flash_attention_class", + current_impl_class=self.__class__, + call_impl_fn=call_impl_fn, ) + @property + def backend_name(self) -> str: + return self.__class__.__name__ + +############ Base ################### class TEFLBackendBase(ABC): @abstractmethod def is_available(self) -> bool: raise NotImplementedError - def get_flash_attention_class(self) -> Type["FlashAttentionBase"]: - raise NotImplementedError - def get_attention_backend(self, attention_params=None): raise NotImplementedError +##### transformer_engine/pytorch/csrc/extensions/pybind.cpp ##### def quantize( self, tensor: torch.Tensor, quantizer: Any, - output: Optional[torch.Tensor] = None, + output: Optional[Any] = None, noop: Optional[torch.Tensor] = None, ) -> Any: raise NotImplementedError def dequantize( self, - input: torch.Tensor, - otype: torch.dtype, - ) -> torch.Tensor: + input: Any, + otype: DType, + ) -> Any: raise NotImplementedError def bgrad_quantize( self, input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError def generic_gemm( self, - A: torch.Tensor, + A: Any, transA: bool, - B: torch.Tensor, + B: Any, transB: bool, - D: torch.Tensor, + D: Any, quantizer: Any, - output_dtype: torch.dtype, + output_dtype: Optional[DType], bias: Optional[torch.Tensor], - bias_type: Any, + bias_type: DType, gelu: bool, gelu_in: Optional[torch.Tensor], grad: bool, @@ -237,91 +404,77 @@ def generic_gemm( accumulate: bool, use_split_accumulator: bool, comm_overlap: Optional[Any] = None, - comm_type: Optional[Any] = None, + comm_type: Optional[CommOverlapType] = None, extra_output: Optional[torch.Tensor] = None, bulk_overlap: bool = False, alpha: float = 1.0, beta: Optional[float] = None, - ) -> Any: - raise NotImplementedError - - def te_general_grouped_gemm( - self, - *args, - **kwargs, - ) -> Any: + ) -> List[Any]: raise NotImplementedError + # GELU and variants # def gelu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def geglu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def qgelu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def qgeglu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - + # ReLU and variants # def relu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def reglu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def srelu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def sreglu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - + # SwiGLU and variants # def silu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def swiglu( self, input: torch.Tensor, quantizer: Any, ) -> Any: raise NotImplementedError - def clamped_swiglu( self, input: torch.Tensor, @@ -330,7 +483,7 @@ def clamped_swiglu( alpha: float = 1.702, ) -> Any: raise NotImplementedError - + # Backward of GELU and variants # def dgelu( self, grad: torch.Tensor, @@ -338,7 +491,6 @@ def dgelu( quantizer: Any, ) -> Any: raise NotImplementedError - def dgeglu( self, grad: torch.Tensor, @@ -346,7 +498,6 @@ def dgeglu( quantizer: Any, ) -> Any: raise NotImplementedError - def dqgelu( self, grad: torch.Tensor, @@ -354,7 +505,6 @@ def dqgelu( quantizer: Any, ) -> Any: raise NotImplementedError - def dqgeglu( self, grad: torch.Tensor, @@ -362,7 +512,7 @@ def dqgeglu( quantizer: Any, ) -> Any: raise NotImplementedError - + # Backward of ReLU and variants # def drelu( self, grad: torch.Tensor, @@ -370,7 +520,6 @@ def drelu( quantizer: Any, ) -> Any: raise NotImplementedError - def dreglu( self, grad: torch.Tensor, @@ -378,7 +527,6 @@ def dreglu( quantizer: Any, ) -> Any: raise NotImplementedError - def dsrelu( self, grad: torch.Tensor, @@ -386,7 +534,6 @@ def dsrelu( quantizer: Any, ) -> Any: raise NotImplementedError - def dsreglu( self, grad: torch.Tensor, @@ -394,7 +541,7 @@ def dsreglu( quantizer: Any, ) -> Any: raise NotImplementedError - + # Backward of SiLU and variants # def dsilu( self, grad: torch.Tensor, @@ -402,7 +549,6 @@ def dsilu( quantizer: Any, ) -> Any: raise NotImplementedError - def dswiglu( self, grad: torch.Tensor, @@ -410,7 +556,6 @@ def dswiglu( quantizer: Any, ) -> Any: raise NotImplementedError - def clamped_dswiglu( self, grad: torch.Tensor, @@ -420,103 +565,193 @@ def clamped_dswiglu( alpha: float = 1.702, ) -> Any: raise NotImplementedError - + # DBias + DAct fusions # def dbias_dgelu( self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError - def dbias_dsilu( self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError - def dbias_drelu( self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError - def dbias_dqgelu( self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError - def dbias_dsrelu( self, grad: torch.Tensor, fwd_input: torch.Tensor, quantizer: Any, - ) -> Tuple[torch.Tensor, Any]: + ) -> List[Any]: raise NotImplementedError - - def layernorm_fwd( + # Permutation functions + def moe_permute_fwd( self, input: torch.Tensor, - weight: torch.Tensor, - bias: Optional[torch.Tensor], - eps: float, - ln_out: Optional[torch.Tensor], - quantizer: Any, - otype: torch.dtype, - sm_margin: int, - zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + dtype: DType, + indices: torch.Tensor, + num_out_tokens: int, + workspace: List[torch.Tensor], + max_expanded_token_num: int, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor]]: raise NotImplementedError - - def layernorm_bwd( + def moe_permute_bwd( self, - dy: torch.Tensor, - x: torch.Tensor, - mu: torch.Tensor, - rsigma: torch.Tensor, - gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: raise NotImplementedError - - def rmsnorm_fwd( + def moe_unpermute_fwd( + self, + input: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + num_tokens: int, + topK: int, + ) -> torch.Tensor: + raise NotImplementedError + def moe_unpermute_bwd( + self, + input_bwd: torch.Tensor, + input_fwd: torch.Tensor, + dtype: DType, + row_id_map: torch.Tensor, + prob: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + raise NotImplementedError + # Softmax functions + def scaled_softmax_forward( + self, + input: torch.Tensor, + scale: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_masked_softmax_forward( + self, + input: torch.Tensor, + mask: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_upper_triang_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_upper_triang_masked_softmax_backward( + self, + output_grads_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_aligned_causal_masked_softmax_forward( + self, + input: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + def scaled_aligned_causal_masked_softmax_backward( + self, + output_grad_: torch.Tensor, + softmax_results_: torch.Tensor, + scale_factor: float, + ) -> torch.Tensor: + raise NotImplementedError + # Other granular functions + def layernorm_fwd( self, input: torch.Tensor, weight: torch.Tensor, + bias: Optional[torch.Tensor], eps: float, - ln_out: Optional[torch.Tensor], + ln_out: Any, quantizer: Any, - otype: torch.dtype, + otype: DType, sm_margin: int, zero_centered_gamma: bool, - ) -> Tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor]: + ) -> List[Any]: + raise NotImplementedError + def layernorm_bwd( + self, + dz: torch.Tensor, + x: torch.Tensor, + mu: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: + raise NotImplementedError + def rmsnorm_fwd( + self, + input: Any, + weight: Any, + eps: float, + ln_out: Any, + quantizer: Any, + otype: DType, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: raise NotImplementedError - def rmsnorm_bwd( self, - dy: torch.Tensor, + dz: torch.Tensor, x: torch.Tensor, rsigma: torch.Tensor, gamma: torch.Tensor, - sm_margin: int = 0, - zero_centered_gamma: bool = False, - eps: float = 1e-5, - ) -> Tuple[torch.Tensor, torch.Tensor]: + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: raise NotImplementedError - def rmsnorm_bwd_add( self, - *args, - **kwargs, - ) -> Any: + dz: torch.Tensor, + x: torch.Tensor, + add: torch.Tensor, + rsigma: torch.Tensor, + gamma: torch.Tensor, + sm_margin: int, + zero_centered_gamma: bool, + ) -> List[Any]: raise NotImplementedError def multi_tensor_quantize( @@ -525,7 +760,6 @@ def multi_tensor_quantize( quantizer_list: List[Any], ) -> List[Any]: raise NotImplementedError - def split_quantize( self, tensor: torch.Tensor, @@ -533,177 +767,290 @@ def split_quantize( quantizer_list: List[Any], ) -> List[Any]: raise NotImplementedError - - def moe_permute_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError - - def moe_permute_bwd(self, *args, **kwargs) -> Any: - raise NotImplementedError - - def moe_unpermute_fwd(self, *args, **kwargs) -> Any: - raise NotImplementedError - - def moe_unpermute_bwd(self, *args, **kwargs) -> Any: + def te_general_grouped_gemm( + self, + A: List[Any], + transa: bool, + B: List[Any], + transb: bool, + D: Optional[List[torch.Tensor]], + D_type: DType, + m_splits: List[int], + bias: List[torch.Tensor], + bias_type: DType, + single_output: bool, + pre_gelu_out: List[torch.Tensor], + grad: bool, + workspace: List[torch.Tensor], + workspaceSizes: int, + accumulate: bool, + use_split_accumulator: bool, + math_sm_count: int, + ) -> Optional[List[torch.Tensor]]: raise NotImplementedError - - def scaled_softmax_forward( + def fp8_transpose( self, input: torch.Tensor, - scale: float, + dtype: DType, + out: Optional[torch.Tensor], ) -> torch.Tensor: raise NotImplementedError - - def scaled_softmax_backward( + def swap_first_dims( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + tensor: torch.Tensor, + out: Optional[torch.Tensor], ) -> torch.Tensor: raise NotImplementedError + def get_fused_attn_backend( + self, + is_training: bool, + q_dtype: DType, + kv_dtype: DType, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + p_dropout: float, + num_attn_heads: int, + num_gqa_groups: int, + max_seqlen_q: int, + max_seqlen_kv: int, + head_dim_qk: int, + head_dim_v: int, + window_size_left: int, + window_size_right: int, + return_max_logit: bool, + ) -> NVTE_Fused_Attn_Backend: + raise NotImplementedError - def scaled_masked_softmax_forward( + def compute_amax( self, input: torch.Tensor, - mask: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax: torch.Tensor, + ) -> None: raise NotImplementedError - - def scaled_masked_softmax_backward( + def fused_amax_and_scale_update_after_reduction( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + amax_reduction_buffer: torch.Tensor, + amax_histories: List[torch.Tensor], + scales: List[torch.Tensor], + amax_compute_algo: str, + fp8_dtype: DType, + margin: float, + ) -> None: raise NotImplementedError - - def scaled_upper_triang_masked_softmax_forward( + def fp8_block_scaling_compute_partial_amax( self, - input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + tensor: torch.Tensor, + amax: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + ) -> None: raise NotImplementedError - - def scaled_upper_triang_masked_softmax_backward( + def fp8_block_scaling_partial_cast( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, - ) -> torch.Tensor: + inp: torch.Tensor, + out: torch.Tensor, + scale: torch.Tensor, + h: int, + w: int, + start_offset: int, + block_len: int, + out_dtype: DType, + ) -> None: raise NotImplementedError - - def scaled_aligned_causal_masked_softmax_forward( + def fused_multi_row_padding( self, input: torch.Tensor, - scale: float, - ) -> torch.Tensor: + output: torch.Tensor, + input_row_list: List[int], + padded_input_row_list: List[int], + ) -> None: + raise NotImplementedError + def fused_multi_row_unpadding( + self, + input: torch.Tensor, + output: torch.Tensor, + input_row_list: List[int], + unpadded_input_row_list: List[int], + ) -> None: raise NotImplementedError - def scaled_aligned_causal_masked_softmax_backward( + # attention kernels + def fa_prepare_fwd( self, - output_grad: torch.Tensor, - softmax_output: torch.Tensor, - scale: float, + qkvi: torch.Tensor, ) -> torch.Tensor: raise NotImplementedError - - def get_fused_attn_backend( + def fa_prepare_bwd( self, - *args, - **kwargs, - ) -> int: + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + ) -> torch.Tensor: raise NotImplementedError - def fused_attn_fwd( self, - *args, - **kwargs, - ) -> Any: + max_seqlen_q: int, + max_seqlen_kv: int, + is_training: bool, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + fake_dtype: torch.dtype, + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + page_table_k: Optional[torch.Tensor], + page_table_v: Optional[torch.Tensor], + s_quantizer: Any, + o_quantizer: Any, + Bias: Optional[torch.Tensor], + SoftmaxOffset: Optional[torch.Tensor], + rng_gen: Optional[torch.Generator], + rng_elts_per_thread: int, + return_max_logit: bool, + ) -> List[Any]: raise NotImplementedError - def fused_attn_bwd( self, - *args, - **kwargs, - ) -> Any: - raise NotImplementedError - - def fa_prepare_fwd( - self, - *args, - **kwargs, - ) -> Any: - raise NotImplementedError - - def fa_prepare_bwd( - self, - *args, - **kwargs, - ) -> Any: + max_seqlen_q: int, + max_seqlen_kv: int, + attn_scale: float, + p_dropout: float, + set_zero: bool, + qkv_layout: NVTE_QKV_Layout, + bias_type: NVTE_Bias_Type, + attn_mask_type: NVTE_Mask_Type, + softmax_type: NVTE_Softmax_Type, + window_size: List[int], + deterministic: bool, + cu_seqlens_q: torch.Tensor, + cu_seqlens_kv: torch.Tensor, + Q: Any, + K: Any, + V: Any, + O: Any, + dO: Any, + fake_dtype: torch.dtype, + dqkv_type: DType, + Aux_CTX_Tensors: List[torch.Tensor], + cu_seqlens_q_padded: Optional[torch.Tensor], + cu_seqlens_kv_padded: Optional[torch.Tensor], + s_quantizer: Any, + dp_quantizer: Any, + dqkv_quantizer: Any, + ) -> List[Any]: raise NotImplementedError - def copy_to_kv_cache( self, - *args, - **kwargs, - ) -> Any: + new_k: torch.Tensor, + new_v: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + page_table: torch.Tensor, + cu_new_lens: torch.Tensor, + cu_cached_lens: torch.Tensor, + qkv_format: NVTE_QKV_Format, + b: int, + max_ctx_len: int, + max_seq_len: int, + max_pages_per_seq: int, + is_non_paged: bool, + ) -> None: raise NotImplementedError - def convert_thd_to_bshd( self, - *args, - **kwargs, - ) -> Any: + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + b: int, + max_seq_len: int, + ) -> torch.Tensor: raise NotImplementedError - def convert_bshd_to_thd( self, - *args, - **kwargs, - ) -> Any: + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + t: int, + ) -> torch.Tensor: raise NotImplementedError + # fused apply rope def fused_rope_forward( self, - *args, - **kwargs, - ) -> Any: + 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: raise NotImplementedError - def fused_rope_backward( self, - *args, - **kwargs, - ) -> Any: + output_grads: torch.Tensor, + freqs: torch.Tensor, + qkv_format: NVTE_QKV_Format, + interleaved: bool, + cu_seqlens: Optional[torch.Tensor], + cp_size: int, + cp_rank: int, + ) -> torch.Tensor: raise NotImplementedError - def fused_qkv_rope_forward( self, - *args, - **kwargs, - ) -> Any: + 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]: raise NotImplementedError - def fused_qkv_rope_backward( self, - *args, - **kwargs, - ) -> Any: + 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: raise NotImplementedError + # fused router def fused_topk_with_score_function_fwd( self, logits: torch.Tensor, topk: int, use_pre_softmax: bool, - num_groups: int, - group_topk: int, - scaling_factor: float, - score_function: Any, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: Optional[float], + score_function: str, expert_bias: Optional[torch.Tensor], - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: raise NotImplementedError - def fused_topk_with_score_function_bwd( self, num_tokens: int, @@ -713,19 +1060,17 @@ def fused_topk_with_score_function_bwd( grad_probs: torch.Tensor, topk: int, use_pre_softmax: bool, - scaling_factor: float, - score_function: Any, - ) -> Any: + scaling_factor: Optional[float], + score_function: str, + ) -> torch.Tensor: raise NotImplementedError - def fused_score_for_moe_aux_loss_fwd( self, logits: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: raise NotImplementedError - def fused_score_for_moe_aux_loss_bwd( self, num_tokens: int, @@ -733,10 +1078,9 @@ def fused_score_for_moe_aux_loss_bwd( intermediate_output: torch.Tensor, grad_scores: torch.Tensor, topk: int, - score_function: Any, - ) -> Any: + score_function: str, + ) -> torch.Tensor: raise NotImplementedError - def fused_moe_aux_loss_fwd( self, probs: torch.Tensor, @@ -747,9 +1091,8 @@ def fused_moe_aux_loss_fwd( num_cols: int, topk: int, coeff: float, - ) -> Any: + ) -> Tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError - def fused_moe_aux_loss_bwd( self, Const_buf: torch.Tensor, @@ -757,177 +1100,117 @@ def fused_moe_aux_loss_bwd( num_rows: int, num_cols: int, grad_aux_loss: torch.Tensor, - ) -> Any: + ) -> torch.Tensor: raise NotImplementedError + # Dropout def dropout_fwd( self, input: torch.Tensor, dropout_probability: float, - out: Optional[torch.Tensor] = None, + out: Optional[torch.Tensor], ) -> Tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError - def dropout_bwd( self, grad_output: torch.Tensor, mask: torch.Tensor, dropout_probability: float, - grad_input: Optional[torch.Tensor] = None, + grad_input: Optional[torch.Tensor], ) -> torch.Tensor: raise NotImplementedError - def fp8_transpose( - self, - input: torch.Tensor, - dtype: Any, - *, - out: torch.Tensor, - ) -> None: - raise NotImplementedError - - def swap_first_dims( - self, - tensor: torch.Tensor, - *, - out: torch.Tensor, - ) -> None: - raise NotImplementedError - - def compute_amax( - self, - input: torch.Tensor, - amax: torch.Tensor, - ) -> None: - raise NotImplementedError - - def fused_amax_and_scale_update_after_reduction( - self, - *args, - **kwargs, - ) -> None: - raise NotImplementedError - - def fp8_block_scaling_compute_partial_amax( - self, - tensor: torch.Tensor, - amax: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - ) -> None: - raise NotImplementedError - - def fp8_block_scaling_partial_cast( - self, - inp: torch.Tensor, - out: torch.Tensor, - scale: torch.Tensor, - h: int, - w: int, - start_offset: int, - block_len: int, - out_dtype: Any, - ) -> None: - raise NotImplementedError - - def fused_multi_row_padding( - self, - *args, - **kwargs, - ) -> Any: - raise NotImplementedError - - def fused_multi_row_unpadding( - self, - *args, - **kwargs, - ) -> Any: - raise NotImplementedError - + # Misc def get_cublasLt_version(self) -> int: raise NotImplementedError - def get_cudnn_version(self) -> int: raise NotImplementedError - def get_num_cublas_streams(self) -> int: raise NotImplementedError + # Support THD format for Context Parallel def thd_read_half_tensor( self, - *args, - **kwargs, - ) -> Any: + tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + half_idx: int, + ) -> torch.Tensor: raise NotImplementedError - def thd_second_half_lse_correction( self, - *args, - **kwargs, - ) -> Any: + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + ) -> None: raise NotImplementedError - def thd_read_second_half_lse( self, - *args, - **kwargs, - ) -> Any: + lse: torch.Tensor, + cu_seqlens: torch.Tensor, + lse_packed: bool, + second_half_lse_seqlen: int, + ) -> torch.Tensor: raise NotImplementedError - def thd_out_correction( self, - *args, - **kwargs, - ) -> Any: + out: torch.Tensor, + out_per_step: torch.Tensor, + lse: torch.Tensor, + lse_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + only_second_half: bool, + lse_packed: bool, + ) -> None: raise NotImplementedError - def thd_grad_correction( self, - *args, - **kwargs, - ) -> Any: + grad: torch.Tensor, + grad_per_step: torch.Tensor, + cu_seqlens: torch.Tensor, + first_half: str, + second_half: str, + ) -> None: raise NotImplementedError - def thd_get_partitioned_indices( self, - *args, - **kwargs, - ) -> Any: + cu_seqlens: torch.Tensor, + total_tokens: int, + world_size: int, + rank: int, + ) -> torch.Tensor: raise NotImplementedError + # nvshmem functions def init_nvshmem_backend( self, - *args, - **kwargs, + process_group: Any, ) -> None: raise NotImplementedError - def create_nvshmem_tensor( self, - *args, - **kwargs, + shape: List[int], + dtype: torch.dtype, ) -> torch.Tensor: raise NotImplementedError - def nvshmem_send_on_current_stream( self, - *args, - **kwargs, + src: torch.Tensor, + dst: torch.Tensor, + peer: int, + signal: torch.Tensor, ) -> None: raise NotImplementedError - def nvshmem_wait_on_current_stream( self, - *args, - **kwargs, + signal: torch.Tensor, + wait_kind: str, ) -> None: raise NotImplementedError - def nvshmem_finalize(self) -> None: raise NotImplementedError + # multi-tensor functions def multi_tensor_scale( self, chunk_size: int, @@ -936,102 +1219,150 @@ def multi_tensor_scale( scale: float, ) -> None: raise NotImplementedError - def multi_tensor_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError - def multi_tensor_unscale_l2norm( self, chunk_size: int, noop_flag: torch.Tensor, tensor_lists: List[List[torch.Tensor]], - scale: torch.Tensor, - per_tensor: bool = False, - ) -> Union[torch.Tensor, List[torch.Tensor]]: + inv_scale: torch.Tensor, + per_tensor: Optional[bool] = False, + ) -> Tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError - def multi_tensor_adam( self, - chunk_size: int = None, - noop_flag: torch.Tensor = None, - tensor_lists: List[List[torch.Tensor]] = None, - lr: float = None, - beta1: float = None, - beta2: float = None, - eps: float = None, - step: int = None, - mode: int = None, - bias_correction: int = None, - weight_decay: float = None, - ): + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + ) -> None: raise NotImplementedError - def multi_tensor_adam_param_remainder( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, ) -> None: raise NotImplementedError - def multi_tensor_adam_fp8( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: float, + beta1: float, + beta2: float, + epsilon: float, + step: int, + mode: int, + bias_correction: int, + weight_decay: float, + fp8_dtype: DType, ) -> None: raise NotImplementedError - def multi_tensor_adam_capturable( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, ) -> None: raise NotImplementedError - def multi_tensor_adam_capturable_master( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + lr: torch.Tensor, + beta1: float, + beta2: float, + epsilon: float, + step: torch.Tensor, + mode: int, + bias_correction: int, + weight_decay: float, + inv_scale: torch.Tensor, ) -> None: raise NotImplementedError - def multi_tensor_sgd( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + wd: float, + momentum: float, + dampening: float, + lr: float, + nesterov: bool, + first_run: bool, + wd_after_momentum: bool, + scale: float, ) -> None: raise NotImplementedError - def multi_tensor_compute_scale_and_scale_inv( self, - *args, - **kwargs, + chunk_size: int, + noop_flag: torch.Tensor, + tensor_lists: List[List[torch.Tensor]], + max_fp8: float, + force_pow_2_scales: bool, + epsilon: float, ) -> None: raise NotImplementedError + # Comm+GEMM Overlap def bulk_overlap_ag_with_external_gemm( self, - allgather_communicator: Any, + allgather_communicator: CommOverlap, send_stream: Any, recv_stream: Any, ) -> Any: raise NotImplementedError +############## class func ################################# def create_fp8_tensor_meta(self) -> FP8TensorMeta: + """Create FP8TensorMeta instance.""" raise NotImplementedError - def create_comm_overlap_helper( self, world_group: Optional[Any] = None, intra_node_group: Optional[Any] = None, - ) -> Any: + ) -> "CommOverlapHelper": + """ + Internal method to create CommOverlapHelper. + Users should use CommOverlapHelper(...) directly. + """ raise NotImplementedError - def create_comm_overlap( self, buffer_shape: List[int], @@ -1047,9 +1378,12 @@ def create_comm_overlap( set_sm_margin: bool = True, atomic_gemm: bool = False, rs_overlap_first_gemm: bool = False, - ) -> Any: + ) -> "CommOverlap": + """ + Internal method to create CommOverlap. + Users should use CommOverlap(...) directly. + """ raise NotImplementedError - def create_comm_overlap_p2p( self, buffer_shape: List[int], @@ -1066,187 +1400,16 @@ def create_comm_overlap_p2p( atomic_gemm: bool = False, use_ce: bool = True, aggregate: bool = False, - ) -> Any: - raise NotImplementedError - -class FlashAttentionBase(torch.nn.Module, ABC): - - def __init__( - self, - softmax_scale: float, - attention_dropout: float = 0.0, - attention_dropout_ctx: Optional[Callable] = None, - attention_type: str = "self", - layer_number: Optional[int] = None, - deterministic: bool = False, - ) -> None: - super().__init__() - - self.softmax_scale = softmax_scale - self.attention_dropout = attention_dropout - self.attention_dropout_ctx = attention_dropout_ctx or nullcontext - self.attention_type = attention_type - self.layer_number = 1 if layer_number is None else layer_number - self.deterministic = deterministic - - # For fallback support - self._manager = None - self._init_params = None - - @abstractmethod - def _forward_impl( - self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, - qkv_layout: str = "sbh3d", - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_kv: Optional[torch.Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_kv: Optional[int] = None, - attn_mask_type: str = "causal", - window_size: Optional[Tuple[int, int]] = None, - alibi_slopes: Optional[torch.Tensor] = None, - cp_group: Optional[Any] = None, - cp_global_ranks: Optional[List[int]] = None, - cp_stream: Optional[torch.cuda.Stream] = None, - cp_comm_type: str = "p2p", - fp8: bool = False, - fp8_meta: Optional[Dict[str, Any]] = None, - quantizers: Optional[Any] = None, - inference_params: Optional[Any] = None, - flash_attention_backend: Optional[Any] = None, - fp8_output: bool = False, - ) -> torch.Tensor: - """ - Actual forward implementation - subclasses must implement this. - - This method contains the backend-specific logic for flash attention. - """ - raise NotImplementedError("Subclasses must implement _forward_impl()") - - def forward( - self, - query_layer: torch.Tensor, - key_layer: torch.Tensor, - value_layer: torch.Tensor, - attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, - qkv_layout: str = "sbh3d", - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_kv: Optional[torch.Tensor] = None, - max_seqlen_q: Optional[int] = None, - max_seqlen_kv: Optional[int] = None, - attn_mask_type: str = "causal", - window_size: Optional[Tuple[int, int]] = None, - alibi_slopes: Optional[torch.Tensor] = None, - cp_group: Optional[Any] = None, - cp_global_ranks: Optional[List[int]] = None, - cp_stream: Optional[torch.cuda.Stream] = None, - cp_comm_type: str = "p2p", - fp8: bool = False, - fp8_meta: Optional[Dict[str, Any]] = None, - quantizers: Optional[Any] = None, - inference_params: Optional[Any] = None, - flash_attention_backend: Optional[Any] = None, - fp8_output: bool = False, - ) -> torch.Tensor: + ) -> "CommOverlapP2P": """ - Forward pass with automatic fallback support and caching. - Delegates to OpManager.call_with_custom_impl for unified dispatch. + Internal method to create CommOverlapP2P. + Users should use CommOverlapP2P(...) directly. """ - if self._manager is None: - return self._forward_impl( - query_layer=query_layer, - key_layer=key_layer, - value_layer=value_layer, - attention_mask=attention_mask, - qkv_layout=qkv_layout, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - attn_mask_type=attn_mask_type, - window_size=window_size, - alibi_slopes=alibi_slopes, - cp_group=cp_group, - cp_global_ranks=cp_global_ranks, - cp_stream=cp_stream, - cp_comm_type=cp_comm_type, - fp8=fp8, - fp8_meta=fp8_meta, - quantizers=quantizers, - inference_params=inference_params, - flash_attention_backend=flash_attention_backend, - fp8_output=fp8_output, - ) - - def call_impl_fn(impl_class): - if impl_class == self.__class__: - return self._forward_impl( - query_layer=query_layer, - key_layer=key_layer, - value_layer=value_layer, - attention_mask=attention_mask, - qkv_layout=qkv_layout, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - attn_mask_type=attn_mask_type, - window_size=window_size, - alibi_slopes=alibi_slopes, - cp_group=cp_group, - cp_global_ranks=cp_global_ranks, - cp_stream=cp_stream, - cp_comm_type=cp_comm_type, - fp8=fp8, - fp8_meta=fp8_meta, - quantizers=quantizers, - inference_params=inference_params, - flash_attention_backend=flash_attention_backend, - fp8_output=fp8_output, - ) - else: - fallback_instance = impl_class(**self._init_params) - fallback_instance._manager = self._manager - fallback_instance._init_params = self._init_params - return fallback_instance._forward_impl( - query_layer=query_layer, - key_layer=key_layer, - value_layer=value_layer, - attention_mask=attention_mask, - qkv_layout=qkv_layout, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - attn_mask_type=attn_mask_type, - window_size=window_size, - alibi_slopes=alibi_slopes, - cp_group=cp_group, - cp_global_ranks=cp_global_ranks, - cp_stream=cp_stream, - cp_comm_type=cp_comm_type, - fp8=fp8, - fp8_meta=fp8_meta, - quantizers=quantizers, - inference_params=inference_params, - flash_attention_backend=flash_attention_backend, - fp8_output=fp8_output, - ) - - return self._manager.call_with_custom_impl( - op_name="get_flash_attention_class", - current_impl_class=self.__class__, - call_impl_fn=call_impl_fn, - ) - - @property - def backend_name(self) -> str: - return self.__class__.__name__ - + raise NotImplementedError + def get_flash_attention_class(self) -> Type["FlashAttentionBase"]: + raise NotImplementedError +############ Wapper ################# class TEFLModule: def __init__(self, manager=None): """ @@ -1259,12 +1422,11 @@ def __init__(self, manager=None): # Import here to avoid circular dependency from .manager import get_default_manager self._manager = manager if manager is not None else get_default_manager() - + # emum self.DType = DType self.Float8BlockScaleTensorFormat = Float8BlockScaleTensorFormat self.FP8FwdTensors = FP8FwdTensors self.FP8BwdTensors = FP8BwdTensors - self.FP8TensorMeta = FP8TensorMeta self.NVTE_Activation_Type = NVTE_Activation_Type self.NVTE_Bias_Type = NVTE_Bias_Type self.NVTE_Mask_Type = NVTE_Mask_Type @@ -1275,14 +1437,11 @@ def __init__(self, manager=None): self.CommOverlapType = CommOverlapType self.CommOverlapAlgo = CommOverlapAlgo self.CommGemmOverlapRole = CommGemmOverlapRole - + # class + self.FP8TensorMeta = FP8TensorMeta self.CommOverlapHelper = CommOverlapHelper self.CommOverlap = CommOverlap self.CommOverlapP2P = CommOverlapP2P - self.CommGemmOverlapAlgoConfig = CommGemmOverlapAlgoConfig - - self.FusedAdamCUDAKernel = FusedAdamCUDAKernel - self.FusedSGDCUDAKernel = FusedSGDCUDAKernel def __getattr__(self, name: str) -> Any: """ @@ -1316,8 +1475,7 @@ def __dir__(self): 'FP8TensorMeta', 'NVTE_Activation_Type', 'NVTE_Bias_Type', 'NVTE_Mask_Type', 'NVTE_Softmax_Type', 'NVTE_Fused_Attn_Backend', 'NVTE_QKV_Format', 'NVTE_QKV_Layout', 'CommOverlapType', 'CommOverlapAlgo', 'CommGemmOverlapRole', - 'CommOverlapHelper', 'CommOverlap', 'CommOverlapP2P', 'CommGemmOverlapAlgoConfig', - 'FusedAdamCUDAKernel', 'FusedSGDCUDAKernel' + 'CommOverlapHelper', 'CommOverlap', 'CommOverlapP2P', ] # Add operator names from OpManager's registry diff --git a/transformer_engine/plugin/tests/test_normalization.py b/transformer_engine/plugin/tests/test_normalization.py index 6a6114a398..1083c8b02c 100644 --- a/transformer_engine/plugin/tests/test_normalization.py +++ b/transformer_engine/plugin/tests/test_normalization.py @@ -13,6 +13,7 @@ TestCase, generate_random_tensor, ) +from transformer_engine.plugin.core.ops import DType class NormalizationTests(TestCase): @@ -57,7 +58,7 @@ def test_layernorm_forward(self, shape=(2, 4, 8)): try: output, mean, rsigma = backend.layernorm_fwd( x, weight, bias, self.eps, - None, None, torch.float32, 0, False + None, None, DType.kFloat32, 0, False ) self.assert_close( output, ref_output, rtol=1e-5, atol=1e-7, @@ -143,7 +144,7 @@ def test_rmsnorm_forward(self, shape=(2, 4, 8)): try: output, _, rsigma = backend.rmsnorm_fwd( x, weight, self.eps, - None, None, torch.float32, 0, False + None, None, DType.kFloat32, 0, False ) self.assert_close( output, ref_output, rtol=1e-5, atol=1e-7, @@ -185,7 +186,7 @@ def test_rmsnorm_backward(self, shape=(2, 4, 8)): grad_x, grad_weight = backend.rmsnorm_bwd( grad_output, x_copy, rsigma.detach(), - weight_copy, 0, False, self.eps + weight_copy, 0, False ) self.assert_close( diff --git a/transformer_engine/plugin/tests/test_operations.py b/transformer_engine/plugin/tests/test_operations.py index 0d64c7e753..0ebe470e91 100644 --- a/transformer_engine/plugin/tests/test_operations.py +++ b/transformer_engine/plugin/tests/test_operations.py @@ -13,6 +13,7 @@ TestCase, generate_random_tensor, ) +from transformer_engine.plugin.core.ops import DType class OperationsTests(TestCase): @@ -39,7 +40,7 @@ def test_gemm_basic(self, M=32, N=64, K=48): output, _, _, _ = backend.generic_gemm( A, False, B, False, D, - None, torch.float32, None, None, + None, DType.kFloat32, None, DType.kFloat32, False, None, False, workspace, 1024, False, False ) @@ -71,7 +72,7 @@ def test_gemm_transpose_a(self, M=32, N=64, K=48): output, _, _, _ = backend.generic_gemm( A, True, B, False, D, - None, torch.float32, None, None, + None, DType.kFloat32, None, DType.kFloat32, False, None, False, workspace, 1024, False, False ) @@ -103,7 +104,7 @@ def test_gemm_3d(self, B=2, M=16, N=32, K=24): output, _, _, _ = backend.generic_gemm( B_mat, False, A, False, D, - None, torch.float32, None, None, + None, DType.kFloat32, None, DType.kFloat32, False, None, False, workspace, 1024, False, False ) @@ -181,7 +182,7 @@ def test_dropout(self, shape=(4, 8, 16)): for backend_name in self.backends: backend = get_backend(backend_name) try: - output, mask = backend.dropout_fwd(x, dropout_prob) + output, mask = backend.dropout_fwd(x, dropout_prob, None) num_nonzero = (output != 0).sum().item() total_elements = output.numel() @@ -206,7 +207,7 @@ def test_dropout(self, shape=(4, 8, 16)): ) grad_output = generate_random_tensor(shape, dtype=torch.bfloat16, device=self.device) - grad_input = backend.dropout_bwd(grad_output, mask, dropout_prob) + grad_input = backend.dropout_bwd(grad_output, mask, dropout_prob, None) grad_nonzero_mask = (grad_input != 0) output_nonzero_mask = (output != 0) diff --git a/transformer_engine/plugin/tests/test_optimizer.py b/transformer_engine/plugin/tests/test_optimizer.py index d4f72919ef..905c7ebbe2 100644 --- a/transformer_engine/plugin/tests/test_optimizer.py +++ b/transformer_engine/plugin/tests/test_optimizer.py @@ -201,7 +201,7 @@ def test_multi_tensor_adam(self, num_tensors=3, shape=(32, 64)): lr=lr, beta1=beta1, beta2=beta2, - eps=eps, + epsilon=eps, step=step, mode=1, # AdamW mode bias_correction=1, @@ -222,6 +222,155 @@ def test_multi_tensor_adam(self, num_tensors=3, shape=(32, 64)): self.failed += 1 print(f" ✗ {backend_name}: {e}") + def _fp32_to_param_remainder(self, fp32_tensor): + """Split FP32 tensor into int16 param (high 16 bits) + int16 remainder (low 16 bits). + + Matches the CUDA split convention: + 1. Extract high 16 bits as param, low 16 bits as remainder. + 2. If remainder < 0, increment param (round up). + """ + int32 = fp32_tensor.view(torch.int32) + rem = (int32 & 0xFFFF).to(torch.int16) + high = ((int32 >> 16) & 0xFFFF).to(torch.int16) + high = torch.where(rem < 0, high + 1, high) + # param is stored as bf16 (same bits as high int16) + param = high.view(torch.bfloat16) + return param, rem + + def _param_remainder_to_fp32(self, param, remainder): + """Reconstruct FP32 from int16 param (high bits) + int16 remainder (low bits). + + Matches the CUDA reconstruct convention: + 1. If remainder < 0, decrement param (undo rounding). + 2. Combine high and low 16 bits into FP32. + """ + local_p = param.view(torch.int16).clone() + local_rem = remainder.clone() + local_p = torch.where(local_rem < 0, local_p - 1, local_p) + high = local_p.to(torch.int32) << 16 + low = local_rem.to(torch.int32) & 0xFFFF + return (high | low).view(torch.float32) + + def _reference_adam_param_remainder( + self, grads, params, exp_avgs, exp_avg_sqs, param_remainders, + lr, beta1, beta2, epsilon, step, mode, bias_correction, weight_decay + ): + """Pure-PyTorch reference for multi_tensor_adam_param_remainder.""" + bc1 = 1 - beta1 ** step if bias_correction else 1.0 + bc2 = 1 - beta2 ** step if bias_correction else 1.0 + is_adamw = (mode == 1) + + for g, p, m, v, p_rem in zip( + grads, params, exp_avgs, exp_avg_sqs, param_remainders + ): + g_float = g.float() + param_master = self._param_remainder_to_fp32(p, p_rem) + + if not is_adamw and weight_decay != 0: + g_float = g_float + weight_decay * param_master + + m.mul_(beta1).add_(g_float, alpha=1 - beta1) + v.mul_(beta2).addcmul_(g_float, g_float, value=1 - beta2) + + m_corr = m / bc1 + v_corr = v / bc2 + denom = torch.sqrt(v_corr) + epsilon + update = m_corr / denom + + if is_adamw and weight_decay != 0: + update = update + weight_decay * param_master + + param_master = param_master - lr * update + + new_p, new_rem = self._fp32_to_param_remainder(param_master) + p.view(torch.int16).copy_(new_p.view(torch.int16)) + p_rem.copy_(new_rem) + + def test_multi_tensor_adam_param_remainder(self, num_tensors=3, shape=(32, 64)): + print(f"\n Testing multi_tensor_adam_param_remainder with {num_tensors} tensors of shape {shape}") + + lr = 0.001 + beta1 = 0.9 + beta2 = 0.999 + eps = 1e-8 + step = 1 + weight_decay = 0.01 + mode = 1 # AdamW + + for backend_name in self.backends: + backend = get_backend(backend_name) + try: + # Create FP32 master weights, then split into param + remainder + master_weights = [generate_random_tensor(shape, dtype=torch.float32, device=self.device) + for _ in range(num_tensors)] + grads = [generate_random_tensor(shape, dtype=torch.bfloat16, device=self.device) + for _ in range(num_tensors)] + + params = [] + remainders = [] + for mw in master_weights: + p, r = self._fp32_to_param_remainder(mw) + params.append(p.clone()) + remainders.append(r.clone()) + + exp_avgs = [torch.zeros(shape, dtype=torch.float32, device=self.device) + for _ in range(num_tensors)] + exp_avg_sqs = [torch.zeros(shape, dtype=torch.float32, device=self.device) + for _ in range(num_tensors)] + + # Clone for reference + ref_params = [p.clone() for p in params] + ref_remainders = [r.clone() for r in remainders] + ref_exp_avgs = [torch.zeros_like(m) for m in exp_avgs] + ref_exp_avg_sqs = [torch.zeros_like(v) for v in exp_avg_sqs] + ref_grads = [g.clone() for g in grads] + + # Reference step + self._reference_adam_param_remainder( + ref_grads, ref_params, ref_exp_avgs, ref_exp_avg_sqs, ref_remainders, + lr, beta1, beta2, eps, step, mode, 1, weight_decay, + ) + + # Backend step + noop_flag = torch.tensor([0], dtype=torch.int32, device=self.device) + backend.multi_tensor_adam_param_remainder( + chunk_size=2048, + noop_flag=noop_flag, + tensor_lists=[grads, params, exp_avgs, exp_avg_sqs, remainders], + lr=lr, + beta1=beta1, + beta2=beta2, + epsilon=eps, + step=step, + mode=mode, + bias_correction=1, + weight_decay=weight_decay, + ) + + # Compare reconstructed FP32 master weights + for i in range(num_tensors): + out_fp32 = self._param_remainder_to_fp32(params[i], remainders[i]) + ref_fp32 = self._param_remainder_to_fp32(ref_params[i], ref_remainders[i]) + self.assert_close( + out_fp32, ref_fp32, rtol=1e-5, atol=1e-7, + msg=f"multi_tensor_adam_param_remainder param {i} mismatch for {backend_name}" + ) + self.assert_close( + exp_avgs[i], ref_exp_avgs[i], rtol=1e-5, atol=1e-7, + msg=f"multi_tensor_adam_param_remainder exp_avg {i} mismatch for {backend_name}" + ) + self.assert_close( + exp_avg_sqs[i], ref_exp_avg_sqs[i], rtol=1e-5, atol=1e-7, + msg=f"multi_tensor_adam_param_remainder exp_avg_sq {i} mismatch for {backend_name}" + ) + print(f" ✓ {backend_name}") + except NotImplementedError: + self.skipped += 1 + print(f" ⊘ {backend_name} (not implemented)") + except Exception as e: + self.failed += 1 + print(f" ✗ {backend_name}: {e}") + def _reference_multi_tensor_unscale_l2norm(self, tensors, inv_scale, per_tensor=False): """Reference implementation for multi_tensor_unscale_l2norm. @@ -258,7 +407,7 @@ def test_multi_tensor_unscale_l2norm(self, num_tensors=4, shape=(64, 128)): chunk_size=2048, noop_flag=noop_flag, tensor_lists=[tensors], - scale=inv_scale, + inv_scale=inv_scale, per_tensor=False ) @@ -298,6 +447,9 @@ def run_all_tests(self): # multi_tensor_adam tests self.test_multi_tensor_adam(num_tensors=3, shape=(32, 64)) + # multi_tensor_adam_param_remainder tests + self.test_multi_tensor_adam_param_remainder(num_tensors=3, shape=(32, 64)) + return self.report() diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index 6c0f969e47..1ca1855f8f 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -508,7 +508,6 @@ def forward( FP8GlobalStateManager.IS_FIRST_FP8_MODULE = _first_fp8_module ctx.wgrad_store = wgrad_store ctx.debug = debug - ctx.eps = eps # ------------------------------------------------------ # Cached state for backward pass is ready... @@ -972,7 +971,6 @@ def wgrad_gemm( ln_weight, ctx.bwd_ln_sm_margin, ctx.zero_centered_gamma, - ctx.eps, ) dgrad = dgrad.reshape(inputmat.size()) dbeta = None diff --git a/transformer_engine/pytorch/ops/basic/rmsnorm.py b/transformer_engine/pytorch/ops/basic/rmsnorm.py index 28126fd44f..05597a14fa 100644 --- a/transformer_engine/pytorch/ops/basic/rmsnorm.py +++ b/transformer_engine/pytorch/ops/basic/rmsnorm.py @@ -232,7 +232,6 @@ def op_backward( w, self._sm_margins["backward"], self.zero_centered_gamma, - self.eps, ) # Clear saved tensors if possible diff --git a/transformer_engine/pytorch/optimizers/__init__.py b/transformer_engine/pytorch/optimizers/__init__.py index a19c797dea..e54a17ae78 100644 --- a/transformer_engine/pytorch/optimizers/__init__.py +++ b/transformer_engine/pytorch/optimizers/__init__.py @@ -13,4 +13,4 @@ ) from .fused_adam import FusedAdam from .fused_sgd import FusedSGD -from .multi_tensor_apply import MultiTensorApply, multi_tensor_applier +from .multi_tensor_apply import MultiTensorApply, multi_tensor_applier \ No newline at end of file diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index b2ddd0adf8..18f7e2031a 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -711,7 +711,7 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N self.multi_tensor_adam_param_remainder, tensor_lists ) else: - apply_multi_tensor_adam(self.multi_tensor_adam(), tensor_lists) + apply_multi_tensor_adam(self.multi_tensor_adam, tensor_lists) if len(p_fp8_model) > 0: tensor_lists = [ g_of_fp8_model, @@ -731,14 +731,14 @@ def apply_multi_tensor_adam(adam_func, tensor_lists, inv_scale=None, out_dtype=N m_of_f32_model, v_of_f32_model, ] - apply_multi_tensor_adam(self.multi_tensor_adam(), tensor_lists) + apply_multi_tensor_adam(self.multi_tensor_adam, tensor_lists) else: # self.master_weights=False and self.capturable=False if len(p_f16_model) > 0: tensor_lists = [g_of_f16_model, p_f16_model, m_of_f16_model, v_of_f16_model] - apply_multi_tensor_adam(self.multi_tensor_adam(), tensor_lists) + apply_multi_tensor_adam(self.multi_tensor_adam, tensor_lists) if len(p_f32_model) > 0: tensor_lists = [g_of_f32_model, p_f32_model, m_of_f32_model, v_of_f32_model] - apply_multi_tensor_adam(self.multi_tensor_adam(), tensor_lists) + apply_multi_tensor_adam(self.multi_tensor_adam, tensor_lists) # Scaling for name in ["exp_avg", "exp_avg_sq", "master_param"]: