Uploaded using `kernel-builder`.

89181fc verified 19 days ago

17 kB

	# SPDX-License-Identifier: MIT
	# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.

	import torch
	import triton
	from typing import Optional
	from .._triton_kernels.normalization.norm import (
	_layernorm_kernel,
	_fused_add_layernorm_kernel,
	_quant_layernorm_kernel,
	_quant_fused_add_layernorm_kernel,
	_layernorm_bwd_dx_fused_triton,
	_layernorm_bwd_dwdb_triton,
	_layernorm_bwd_dwdb_triton_v2,
	)
	from ..utils.types import get_dtype_max
	from ..utils.logger import AiterTritonLogger

	_LOGGER = AiterTritonLogger()


	def _layernorm_forward(
	y: torch.Tensor,
	x: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	mean: torch.Tensor,
	rstd: torch.Tensor,
	eps: float = 1e-5,
	):

	M, N = x.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // x.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	_layernorm_kernel[(M,)](
	x, y, weight, bias, mean, rstd, x.stride(0), y.stride(0), M, N, eps, BLOCK_SIZE
	)

	return


	def _layernorm_forward_with_add(
	y: torch.Tensor,
	x: torch.Tensor,
	res_in: torch.Tensor,
	res_out: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	mean: torch.Tensor,
	rstd: torch.Tensor,
	epsilon: float,
	):

	M, N = x.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // x.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	_fused_add_layernorm_kernel[(M,)](
	x,
	y,
	res_in,
	res_out,
	weight,
	bias,
	mean,
	rstd,
	x.stride(0),
	y.stride(0),
	M,
	N,
	epsilon,
	BLOCK_SIZE,
	)

	return


	def _layernorm_backward(
	dy: torch.Tensor,
	dx: torch.Tensor,
	dw: torch.Tensor,
	db: torch.Tensor,
	x: torch.Tensor,
	gamma: torch.Tensor,
	mu: torch.Tensor,
	rsigma: torch.Tensor,
	):

	M, N = x.shape
	# calculate dw and db separately when M is small
	IGNORE_DW_DB_IN_FUSED = M <= 512
	tile_num = max(min(256, M // 4), 1)
	if M <= 512 and M * N < 64 * 1024 * 1024:
	tile_num = M
	elif M >= 8192:
	tile_num = 2048
	max_fused_size = 32768 // x.element_size()
	next_power = triton.next_power_of_2(N)
	BLOCK_SIZE = min(max_fused_size, next_power)
	# For cases with small M and large N, decrease block size to help with occupancy and register spill
	if tile_num == M:
	if tile_num > 256:
	BLOCK_SIZE = min(BLOCK_SIZE, 2048)
	else:
	BLOCK_SIZE = min(BLOCK_SIZE, 4096)
	USE_BLOCKED = N > BLOCK_SIZE
	num_warps = min(max(BLOCK_SIZE // 256, 1), 8)

	if not IGNORE_DW_DB_IN_FUSED:
	_dw = torch.zeros((tile_num, N), dtype=torch.float32, device=gamma.device)
	_db = torch.zeros((tile_num, N), dtype=torch.float32, device=gamma.device)
	else:
	_dw = None
	_db = None

	grid_bwd = (tile_num,)
	_layernorm_bwd_dx_fused_triton[grid_bwd](
	dx,
	dy,
	_dw,
	_db,
	x,
	gamma,
	mu,
	rsigma,
	x.stride(0),
	N,
	NUM_ROWS=M,
	BLOCK_SIZE_N=BLOCK_SIZE,
	USE_BLOCKED=USE_BLOCKED,
	num_warps=num_warps,
	IGNORE_DW_DB=IGNORE_DW_DB_IN_FUSED,
	)
	grid_reduce = lambda meta: (triton.cdiv(N, meta["BLOCK_SIZE_N"]),) # noqa: E731
	if not IGNORE_DW_DB_IN_FUSED:
	dwdb_block_n = max(16, N // 256)
	dwdb_block_n = triton.next_power_of_2(dwdb_block_n)
	dwdb_block_m = (64 * 128) // dwdb_block_n
	dwdb_block_m = min(triton.next_power_of_2(tile_num), dwdb_block_m)
	_layernorm_bwd_dwdb_triton[grid_reduce](
	_dw,
	_db,
	dw,
	db,
	min(tile_num, M),
	N,
	BLOCK_SIZE_M=dwdb_block_m,
	BLOCK_SIZE_N=dwdb_block_n,
	)
	else:
	dwdb_block_n = max(16, N // 256)
	dwdb_block_n = triton.next_power_of_2(dwdb_block_n)
	dwdb_block_m = (64 * 128) // dwdb_block_n
	dwdb_block_m = min(triton.next_power_of_2(M), dwdb_block_m)
	_layernorm_bwd_dwdb_triton_v2[grid_reduce](
	x,
	dy,
	mu,
	rsigma,
	x.stride(0),
	dw,
	db,
	M,
	N,
	BLOCK_SIZE_M=dwdb_block_m,
	BLOCK_SIZE_N=dwdb_block_n,
	)

	return dx, dw, db


	class _LayerNorm(torch.autograd.Function):

	@staticmethod
	def forward(ctx, x, weight, bias, eps, is_grad_enabled):

	is_grad = is_grad_enabled and any(
	tensor.requires_grad for tensor in [x, weight, bias]
	)

	y = torch.empty_like(x)
	M = x.shape[0]
	mean = torch.empty((M,), dtype=torch.float32, device=x.device)
	rstd = torch.empty((M,), dtype=torch.float32, device=x.device)

	_layernorm_forward(y, x, weight, bias, mean, rstd, eps)

	if is_grad:
	ctx.save_for_backward(x, weight, mean, rstd)

	return y

	@staticmethod
	def backward(ctx, dy):
	x, w, m, v = ctx.saved_tensors
	N = w.shape[0]

	dw = torch.empty((N,), dtype=w.dtype, device=w.device)
	db = torch.empty((N,), dtype=w.dtype, device=w.device)
	dx = torch.empty_like(dy)

	_layernorm_backward(dy, dx, dw, db, x, w, m, v)

	return dx, dw, db, None, None


	class _Layernorm2dFwdWithAdd(torch.autograd.Function):

	@staticmethod
	def forward(ctx, y, x, res_in, res_out, weight, bias, eps, is_grad_enabled):

	is_grad = is_grad_enabled and any(
	tensor.requires_grad for tensor in [x, weight, bias]
	)

	M = x.shape[0]
	mean = torch.empty((M,), dtype=torch.float32, device=x.device)
	rstd = torch.empty((M,), dtype=torch.float32, device=x.device)

	_layernorm_forward_with_add(
	y, x, res_in, res_out, weight, bias, mean, rstd, eps
	)

	if is_grad:
	ctx.save_for_backward(res_out, weight, mean, rstd)

	return y

	@staticmethod
	def backward(ctx, dy):
	x, w, m, v = ctx.saved_tensors
	N = w.shape[0]

	dw = torch.empty((N,), dtype=w.dtype, device=w.device)
	db = torch.empty((N,), dtype=w.dtype, device=w.device)
	dx = torch.empty_like(dy)

	_layernorm_backward(dy, dx, dw, db, x, w, m, v)

	return None, dx, None, None, dw, db, None, None


	def layer_norm(
	input: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	eps: float = 1e-5,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Applies Layer Normalization over a mini-batch of inputs.

	Key parameters:
	- input: The input tensor to be normalized with shape (M, N).
	- weight: The learnable weights tensor with shape (N, ).
	- bias: The learnable bias tensor with shape (N, )
	- eps: A value added to the denominator for numerical stability.

	Returns:
	- Output: The output tensor with shape (M, N).
	"""
	_LOGGER.info(f"LAYERNORM: input={tuple(input.shape)} weight={tuple(weight.shape)} ")
	return _LayerNorm.apply(input, weight, bias, eps, torch.is_grad_enabled())


	def layernorm2d_fwd_with_add(
	out: torch.Tensor,
	input: torch.Tensor,
	residual_in: torch.Tensor,
	residual_out: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	epsilon: float,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Adds two inputs and then applies Layer Normalization

	Key parameters:
	- out: The output of layer normalization with shape (M, N). Allocated by the caller
	- input: The input tensor to be normalized with shape (M, N).
	- residual_in: Tensor added to the input and same shape as input (M, N)
	- residual_out: Output tensor that is input + residual_in with shape (M, N). Must be allocated by the caller
	- weight: The learnable weights tensor with shape (N, ).
	- bias: Bias added to the result of layer norm with shape (N,)
	- epsilon: A value added to the denominator for numerical stability.

	Returns:
	- out: The output tensor with shape (M, N).
	- residual_out: Output tensor that is input + residual_in with shape (M, N).
	"""
	_LOGGER.info(
	f"LAYERNORM_2D_FWD_ADD: input={tuple(input.shape)} weight={tuple(weight.shape)} residual_in={tuple(residual_in.shape)} "
	)
	return _Layernorm2dFwdWithAdd.apply(
	out,
	input,
	residual_in,
	residual_out,
	weight,
	bias,
	epsilon,
	torch.is_grad_enabled(),
	)


	def layernorm2d_fwd_with_dynamicquant(
	out: torch.Tensor,
	input: torch.Tensor,
	yscale: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	epsilon: float = 1e-5,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Applies Layer Normalization and then quantizes the output

	Key parameters:
	- out: The output of layer normalization with shape (M, N). Allocated by the caller
	- input: The input tensor to be normalized with shape (M, N) and dtype in (fp32, fp16 or bf16)
	- yscale: Output scale tensor with shape (M,) and dtype fp32. Allocated by the caller
	- weight: The learnable weights tensor with shape (N, ).
	- bias: Bias added to the result of layer norm with shape (N,)
	- eps: A value added to the denominator for numerical stability.

	Returns:
	- out: The output tensor with shape (M, N).
	- yscale: Output scale tensor with shape (M,). Allocated by the caller
	"""
	_LOGGER.info(
	f"LAYERNORM_2D_FWD_DYNAMICQUANT: input={tuple(input.shape)} weight={tuple(weight.shape)} yscale={tuple(yscale.shape)} "
	)
	M, N = input.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // input.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	xscale = None
	IS_SMOOTH = False
	DTYPE_MAX = get_dtype_max(out.dtype)

	# Auxiliary tensor to store the RMSNorm output as fp32 before applying the quantization when using the blocked approach
	aux = torch.empty(M, N, dtype=torch.float32, device=input.device)

	_quant_layernorm_kernel[(M,)](
	input,
	out,
	weight,
	bias,
	xscale,
	yscale,
	aux,
	input.stride(0),
	out.stride(0),
	aux.stride(0),
	M,
	N,
	epsilon,
	DTYPE_MAX,
	IS_SMOOTH,
	BLOCK_SIZE,
	)

	return


	def layernorm2d_fwd_with_smoothquant(
	out: torch.Tensor,
	input: torch.Tensor,
	xscale: torch.Tensor,
	yscale: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	epsilon: float = 1e-5,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Applies Layer Normalization and then quantizes the output

	Key parameters:
	- input: The input tensor to be normalized with shape (M, N).
	- xscale: Input scale tensor which is multiplied with the output of layer normalization before quantization.
	- yscale: Output scale tensor with shape (M,) and dtype fp32. Allocated by the caller
	- weight: The learnable weights tensor with shape (N, ).
	- bias: Bias added to the result of layer norm with shape (N,)
	- eps: A value added to the denominator for numerical stability.

	Returns:
	- Output: The output tensor with shape (M, N).
	"""
	_LOGGER.info(
	f"RMSNORM_2D_FWD_SMOOTHQUANT: input={tuple(input.shape)} weight={tuple(weight.shape)} xscale={tuple(xscale.shape)} yscale={tuple(yscale.shape)} "
	)
	M, N = input.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // input.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	IS_SMOOTH = True
	DTYPE_MAX = get_dtype_max(out.dtype)

	# Auxiliary tensor to store the RMSNorm output as fp32 before applying the quantization when using the blocked approach
	aux = torch.empty(M, N, dtype=torch.float32, device=input.device)

	_quant_layernorm_kernel[(M,)](
	input,
	out,
	weight,
	bias,
	xscale,
	yscale,
	aux,
	input.stride(0),
	out.stride(0),
	aux.stride(0),
	M,
	N,
	epsilon,
	DTYPE_MAX,
	IS_SMOOTH,
	BLOCK_SIZE,
	)

	return


	def layernorm2d_fwd_with_add_dynamicquant(
	out: torch.Tensor,
	input: torch.Tensor,
	residual_in: torch.Tensor,
	residual_out: torch.Tensor,
	yscale: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	epsilon: float = 1e-5,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Adds two input toegether, then does layer Normalization before quantizing the final output

	Key parameters:
	- out: The output of layer normalization with shape (M, N). Allocated by the caller
	- input: The input tensor to be normalized with shape (M, N) and dtype in (fp32, fp16 or bf16)
	- residual_in: Tensor added to the input and same shape as input (M, N)
	- residual_out: Output tensor that is input + residual_in with shape (M, N). Must be allocated by the caller
	- yscale: Output scale tensor with shape (M,) and dtype fp32. Allocated by the caller
	- weight: The learnable weights tensor with shape (N, ).
	- bias: Bias added to the result of layer norm with shape (N,)
	- eps: A value added to the denominator for numerical stability.

	Returns:
	- out: The output tensor with shape (M, N).
	- yscale: Output scale tensor with shape (M,). Allocated by the caller
	"""
	_LOGGER.info(
	f"LAYERNORM_2D_FWD_ADD_DYNAMICQUANT: input={input.shape} weight={weight.shape} residual_in={residual_in.shape} yscale={yscale.shape} "
	)
	M, N = input.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // input.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	xscale = None
	IS_SMOOTH = False
	DTYPE_MAX = get_dtype_max(out.dtype)

	# Auxiliary tensor to store the RMSNorm output as fp32 before applying the quantization when using the blocked approach
	aux = torch.empty(M, N, dtype=torch.float32, device=input.device)

	_quant_fused_add_layernorm_kernel[(M,)](
	input,
	out,
	residual_in,
	residual_out,
	weight,
	bias,
	xscale,
	yscale,
	aux,
	input.stride(0),
	out.stride(0),
	aux.stride(0),
	M,
	N,
	epsilon,
	DTYPE_MAX,
	IS_SMOOTH,
	BLOCK_SIZE,
	)

	return


	def layernorm2d_fwd_with_add_smoothquant(
	out: torch.Tensor,
	input: torch.Tensor,
	residual_in: torch.Tensor,
	residual_out: torch.Tensor,
	xscale: torch.Tensor,
	yscale: torch.Tensor,
	weight: torch.Tensor,
	bias: torch.Tensor,
	epsilon: float = 1e-5,
	x_bias: Optional[torch.Tensor] = None,
	) -> torch.Tensor:
	"""
	Applies Layer Normalization and then quantizes the output

	Key parameters:
	- input: The input tensor to be normalized with shape (M, N).
	- residual_in: Tensor added to the input and same shape as input (M, N)
	- residual_out: Output tensor that is input + residual_in with shape (M, N). Must be allocated by the caller
	- xscale: Input scale tensor which is multiplied with the output of layer normalization before quantization.
	- yscale: Output scale tensor with shape (M,) and dtype fp32. Allocated by the caller
	- weight: The learnable weights tensor with shape (N, ).
	- bias: Bias added to the result of layer norm with shape (N,)
	- eps: A value added to the denominator for numerical stability.

	Returns:
	- Output: The output tensor with shape (M, N).
	- yscale: Output scale tensor with shape (M,). Allocated by the caller
	"""

	_LOGGER.info(
	f"LAYERNORM_2D_FWD_ADD_SMOOTHQUANT: input={tuple(input.shape)} weight={tuple(weight.shape)} "
	+ f"residual_in={tuple(residual_in.shape)} xscale={tuple(xscale.shape)} yscale={tuple(yscale.shape)} "
	)
	M, N = input.shape

	# Less than 64KB per feature: enqueue fused kernel
	MAX_FUSED_SIZE = 65536 // input.element_size()
	BLOCK_SIZE = min(MAX_FUSED_SIZE, triton.next_power_of_2(N))

	IS_SMOOTH = True
	DTYPE_MAX = get_dtype_max(out.dtype)

	# Auxiliary tensor to store the RMSNorm output as fp32 before applying the quantization when using the blocked approach
	aux = torch.empty(M, N, dtype=torch.float32, device=input.device)

	_quant_fused_add_layernorm_kernel[(M,)](
	input,
	out,
	residual_in,
	residual_out,
	weight,
	bias,
	xscale,
	yscale,
	aux,
	input.stride(0),
	out.stride(0),
	aux.stride(0),
	M,
	N,
	epsilon,
	DTYPE_MAX,
	IS_SMOOTH,
	BLOCK_SIZE,
	)

	return