from functools import cache from typing import Optional import torch import triton from .. import _aiter_compat as aiter from .._triton_kernels.quant.fused_fp8_quant import ( _fused_rms_fp8_per_tensor_static_quant_kernel, _fused_rms_fp8_group_quant_kernel, _fused_rms_gated_fp8_group_quant_kernel, _fused_flatten_fp8_group_quant_kernel, _fused_reduce_act_mul_fp8_group_quant, _fused_reduce_rms_fp8_group_quant_kernel, _fused_silu_mul_fp8_per_tensor_static_quant_kernel, ) from ..utils.types import get_fp8_e4m3_dtype from .._triton_kernels.activation import ( _get_activation_from_str, ) from ..utils.logger import AiterTritonLogger _LOGGER = AiterTritonLogger() fp8_dtype = aiter.dtypes.fp8 def fused_rms_fp8_per_tensor_static_quant( inp1, inp1_weight, inp1_epsilon, inp1_scale, inp2=None, inp2_weight=None, inp2_epsilon=None, dtype_quant=fp8_dtype, res1=None, output_unquantized_inp1=False, rmsnorm_convert_to_inp1_type=False, ): """ This op contains several steps: 1. if res1 is not None, inp1 = inp1 + res1, and store inp1 to out_res1 2. perform RMS norm along the last dimenion for inp1 3. if inp2 is not None, perform RMS norm along the last dimenion for inp2 4. perform fp8 quantization for inp1 only Key parameters: - x: Matrix X with shape (M, N1, N2). Returns: - out1_fp8: The output matrix with shape (M, N1). - out1_s: The output matrix with shape (1,). - out1: The output matrix with shape (M, N1). - out2: The output matrix with shape (M, N2). - out_res1: The output matrix with shape (M, N1). - out1: The output matrix with shape (M, N1). """ M, N1 = inp1.shape BLOCK_SIZE_N = triton.next_power_of_2(N1) if inp2 is not None: M2, N2 = inp2.shape BLOCK_SIZE_N = triton.next_power_of_2(N2) assert ( M == M2 ), "The leading dimension should be identical between inp1 and inp2" else: N2 = 0 out1_fp8 = torch.empty((M, N1), dtype=dtype_quant, device=inp1.device) out2 = None out2_row_stride = 0 out2_col_stride = 0 inp2_row_stride = 0 inp2_col_stride = 0 if inp2 is not None: out2 = torch.empty((M, N2), dtype=inp1.dtype, device=inp1.device) inp2_row_stride = inp2.stride(0) inp2_col_stride = inp2.stride(1) out2_row_stride = out2.stride(0) out2_col_stride = out2.stride(1) out1 = None out1_row_stride = 0 out1_col_stride = 0 if output_unquantized_inp1: out1 = torch.empty((M, N1), dtype=inp1.dtype, device=inp1.device) out1_row_stride = out1.stride(0) out1_col_stride = out1.stride(1) out_res1 = None res1_row_stride = 0 res1_col_stride = 0 out_res1_row_stride = 0 out_res1_col_stride = 0 if res1 is not None: Mr, Nr = res1.shape assert ( M == Mr and N1 == Nr ), "The shape should be identical between inp1 and res1" out_res1 = torch.empty((M, N1), dtype=inp1.dtype, device=inp1.device) res1_row_stride = res1.stride(0) res1_col_stride = res1.stride(1) out_res1_row_stride = out_res1.stride(0) out_res1_col_stride = out_res1.stride(1) if BLOCK_SIZE_N <= 512: num_warps = 1 elif BLOCK_SIZE_N <= 2048: num_warps = 4 elif BLOCK_SIZE_N <= 4096: num_warps = 8 else: num_warps = 16 DTYPE_MAX = ( torch.finfo(out1_fp8.dtype).max if torch.is_floating_point(out1_fp8) else torch.iinfo(out1_fp8.dtype).max ) _fused_rms_fp8_per_tensor_static_quant_kernel[(M,)]( inp1, inp1_weight, inp2, inp2_weight, res1, out1_fp8, out2, out_res1, out1, inp1_scale, inp1_epsilon, inp2_epsilon, M, N1, N2, inp1.stride(0), inp2_row_stride, inp1.stride(1), inp2_col_stride, res1_row_stride, res1_col_stride, out1_fp8.stride(0), out1_fp8.stride(1), out2_row_stride, out2_col_stride, out_res1_row_stride, out_res1_col_stride, out1_row_stride, out1_col_stride, BLOCK_SIZE_N=BLOCK_SIZE_N, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, HAVE_SECOND_INPUT=(inp2 is not None), FIRST_INPUT_RES=(res1 is not None), FIRST_INPUT_OUT=output_unquantized_inp1, RMSNORM_CONVERT_TO_INP1_TYPE=rmsnorm_convert_to_inp1_type, num_warps=num_warps, ) return out1_fp8, out1, out2, out_res1 def fused_rms_fp8_group_quant( inp1, inp1_weight, inp1_epsilon, inp2=None, inp2_weight=None, inp2_epsilon=None, group_size=128, dtype_quant=fp8_dtype, res1=None, output_unquantized_inp1=False, transpose_scale=False, ): """ This op contains several steps: 1. if res1 is not None, inp1 = inp1 + res1, and store inp1 to out_res1 2. perform RMS norm along the last dimenion for inp1 3. if inp2 is not None, perform RMS norm along the last dimenion for inp2 4. perform fp8 quantization for inp1 only Key parameters: - x: Matrix X with shape (M, N1, N2). - transpose_scale: If True, return scale with shape (M, cdiv(N1, group_size)) but stored in column-major (transposed) memory layout. Equivalent to: scale.transpose(0, 1).contiguous().view(*scale.shape) Returns: - out1_fp8: The output matrix with shape (M, N1). - out1_bs: The output matrix with shape (M, cdiv(N1, group_size)). When transpose_scale=True, has column-major memory layout (transposed storage). - out1: The output matrix with shape (M, N1). - out2: The output matrix with shape (M, N2). - out_res1: The output matrix with shape (M, N1). - out1: The output matrix with shape (M, N1). """ M, N1 = inp1.shape BLOCK_SIZE_N = max(triton.next_power_of_2(N1), group_size) if inp2 is not None: M2, N2 = inp2.shape BLOCK_SIZE_N = max(triton.next_power_of_2(N2), BLOCK_SIZE_N) assert ( M == M2 ), "The leading dimension should be identical between inp1 and inp2" else: N2 = 0 out1_fp8 = torch.empty((M, N1), dtype=dtype_quant, device=inp1.device) num_bs_cols = (N1 + group_size - 1) // group_size if transpose_scale: # Create with transposed shape for direct transposed storage out1_bs = torch.empty( (num_bs_cols, M), dtype=torch.float32, device=inp1.device, ) else: out1_bs = torch.empty( (M, num_bs_cols), dtype=torch.float32, device=inp1.device, ) out2 = None out2_row_stride = 0 out2_col_stride = 0 inp2_row_stride = 0 inp2_col_stride = 0 if inp2 is not None: out2 = torch.empty((M, N2), dtype=inp1.dtype, device=inp1.device) inp2_row_stride = inp2.stride(0) inp2_col_stride = inp2.stride(1) out2_row_stride = out2.stride(0) out2_col_stride = out2.stride(1) out1 = None out1_row_stride = 0 out1_col_stride = 0 if output_unquantized_inp1: out1 = torch.empty((M, N1), dtype=inp1.dtype, device=inp1.device) out1_row_stride = out1.stride(0) out1_col_stride = out1.stride(1) BLOCK_SIZE_N = max(BLOCK_SIZE_N, group_size) out_res1 = None res1_row_stride = 0 res1_col_stride = 0 out_res1_row_stride = 0 out_res1_col_stride = 0 if res1 is not None: Mr, Nr = res1.shape assert ( M == Mr and N1 == Nr ), "The shape should be identical between inp1 and res1" out_res1 = torch.empty((M, N1), dtype=inp1.dtype, device=inp1.device) res1_row_stride = res1.stride(0) res1_col_stride = res1.stride(1) out_res1_row_stride = out_res1.stride(0) out_res1_col_stride = out_res1.stride(1) if BLOCK_SIZE_N <= 512: num_warps = 1 elif BLOCK_SIZE_N <= 2048: num_warps = 4 elif BLOCK_SIZE_N <= 4096: num_warps = 8 else: num_warps = 16 DTYPE_MAX = ( torch.finfo(out1_fp8.dtype).max if torch.is_floating_point(out1_fp8) else torch.iinfo(out1_fp8.dtype).max ) # When transpose_scale=True, swap the strides to write directly in transposed layout if transpose_scale: out1_bs_row_stride = out1_bs.stride(1) out1_bs_col_stride = out1_bs.stride(0) else: out1_bs_row_stride = out1_bs.stride(0) out1_bs_col_stride = out1_bs.stride(1) _fused_rms_fp8_group_quant_kernel[(M,)]( inp1, inp1_weight, inp2, inp2_weight, res1, out1_fp8, out1_bs, out2, out_res1, out1, inp1_epsilon, inp2_epsilon, M, N1, N2, inp1.stride(0), inp2_row_stride, inp1.stride(1), inp2_col_stride, res1_row_stride, res1_col_stride, out1_fp8.stride(0), out1_fp8.stride(1), out1_bs_row_stride, out1_bs_col_stride, out2_row_stride, out2_col_stride, out_res1_row_stride, out_res1_col_stride, out1_row_stride, out1_col_stride, BLOCK_SIZE_N=BLOCK_SIZE_N, QUANT_BLOCK_SIZE=group_size, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, HAVE_SECOND_INPUT=(inp2 is not None), FIRST_INPUT_RES=(res1 is not None), FIRST_INPUT_OUT=output_unquantized_inp1, num_warps=num_warps, ) # When transpose_scale=True, view the transposed buffer back to original shape # This keeps shape (M, num_bs_cols) but with column-major memory layout if transpose_scale: out1_bs = out1_bs.view(M, num_bs_cols) return (out1_fp8, out1_bs), out1, out2, out_res1 def get_fp8_min_max_bounds(fp8_dtype: torch.dtype) -> tuple[float, float]: """Match vLLM ``quant_utils.get_fp8_min_max`` for ``fp8_dtype`` (incl. ROCm fnuz ±224).""" if fp8_dtype == torch.float8_e4m3fnuz: return -224.0, 224.0 finfo = torch.finfo(fp8_dtype) return float(finfo.min), float(finfo.max) @cache def _num_compute_units(device_id: int = 0) -> int: """Match vLLM ``vllm.utils.platform_utils.num_compute_units`` (``current_platform.num_compute_units``).""" return torch.cuda.get_device_properties(device_id).multi_processor_count def calc_rows_per_block(M: int, device: torch.device) -> int: """Same heuristic as vLLM ``input_quant_fp8.calc_rows_per_block``.""" if device.type != "cuda": raise ValueError( "fused_rms_gated_fp8_group_quant targets AMD ROCm (HIP); expected a CUDA/HIP device." ) device_id = ( device.index if device.index is not None else torch.cuda.current_device() ) sm_count = max(int(_num_compute_units(device_id)), 1) rows_per_block = triton.next_power_of_2(triton.cdiv(M, 2 * sm_count)) return min(int(rows_per_block), 4) def fused_rms_gated_fp8_group_quant( x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None, z: torch.Tensor, eps: float, *, norm_before_gate: bool = True, use_ue8m0: bool = False, activation: str = "silu", out_dtype: torch.dtype | None = None, fp8_min: float | None = None, fp8_max: float | None = None, fp8_min_scaling_factor: float | None = None, group_size: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Fused RMSNorm (with optional bias), optional multiplicative gate from ``z``, and FP8 quantization (same contract as vLLM ``_rmsnorm_quantize_group_native`` for ``group_size == N``). Comparison with ``fused_rms_fp8_group_quant``: Use ``fused_rms_fp8_group_quant`` when you need optional **two-stream** RMSNorm (``inp1`` / optional ``inp2`` with separate weights and epsilons), optional **residual** fused into ``inp1`` (``res1``), FP8 group quantization on the **first** normalized stream only, the richer return tuple (quantized FP8, block scales, optional unquantized ``inp1``, second RMS output, residual output), and optional ``transpose_scale`` layout for scales. Use **this** function for **single** hidden ``x``, one RMS **weight** (and optional **bias**), plus ``z`` for **elementwise multiplicative gating** (SiLU / sigmoid-style activations on ``z``) matching ``x``'s shape; optional ``norm_before_gate`` ordering; vLLM-aligned FP8 bounds / optional UE8M0 / ``group_size`` (``None`` = one scale per row, else per-column-group scales). Returns only ``(x_quant_fp8, scales)``. Suited to gated RMSNorm input quantization (e.g. SwiGLU-style / vLLM ``_rmsnorm_quantize_group_native`` contracts), not the two-stream + residual pattern above. ``x`` and ``z`` must be 2D contiguous with identical shape ``(M, N)``. Returns ``(x_quant_fp8, scales)`` where ``scales`` is ``(M,)`` float32 if ``group_size`` is ``None`` (one scale per row), or ``(M, N // group_size)`` float32 when ``group_size`` divides ``N`` (one scale per row per column group). ``fp8_min`` / ``fp8_max`` / ``fp8_min_scaling_factor`` default from ``out_dtype`` (or ``get_fp8_e4m3_dtype()``) using the same rules as vLLM ``get_fp8_min_max`` and ``1.0 / (_FP8_MAX * 512)``. Pass them explicitly when you want to pin values (e.g. from vLLM's ``get_fp8_min_max()`` at model init). Raises: ValueError: if ``group_size`` is not ``None`` and ``group_size > N``, ``group_size <= 0``, or ``N`` is not divisible by ``group_size``. """ assert x.is_contiguous() and z.is_contiguous() assert x.shape == z.shape, "x and z must have the same shape" fp8_dtype = out_dtype if out_dtype is not None else get_fp8_e4m3_dtype() if (fp8_min is None) ^ (fp8_max is None): raise ValueError("fp8_min and fp8_max must be passed together or both omitted.") if fp8_min is None: fp8_min, fp8_max = get_fp8_min_max_bounds(fp8_dtype) if fp8_min_scaling_factor is None: fp8_min_scaling_factor = 1.0 / (fp8_max * 512.0) weight = weight.contiguous() if bias is not None: bias = bias.contiguous() M, N = x.shape if group_size is not None: if group_size <= 0: raise ValueError(f"group_size must be positive, got {group_size}") if group_size > N: raise ValueError( f"group_size ({group_size}) must be less than or equal to hidden size " f"N ({N}); per-column FP8 groups cannot exceed the row width." ) if N % group_size != 0: raise ValueError( f"hidden size N ({N}) must be divisible by group_size ({group_size})." ) effective_gs = N if group_size is None else int(group_size) num_groups = N // effective_gs MAX_FUSED_SIZE = 65536 // x.element_size() if N > MAX_FUSED_SIZE: raise RuntimeError("This RMSNorm quant kernel does not support N >= 64KB.") rms_tile = min(512, triton.next_power_of_2(N)) block_g = triton.next_power_of_2(effective_gs) rows_per_block = calc_rows_per_block(M, x.device) num_warps = min(max(block_g // 256, 1), 8) x_quant = torch.empty(M, N, dtype=fp8_dtype, device=x.device) if group_size is None: scales = torch.empty(M, dtype=torch.float32, device=x.device) stride_s_row = int(scales.stride(0)) stride_s_g = 0 else: scales = torch.empty(M, num_groups, dtype=torch.float32, device=x.device) stride_s_row, stride_s_g = (int(scales.stride(0)), int(scales.stride(1))) grid = (triton.cdiv(M, rows_per_block),) _fused_rms_gated_fp8_group_quant_kernel[grid]( x, weight, bias, z, x_quant, scales, x.stride(0), z.stride(0), x_quant.stride(0), stride_s_row, stride_s_g, M, N, eps, RMS_TILE=rms_tile, ROWS_PER_BLOCK=rows_per_block, GROUP_SIZE=effective_gs, NUM_GROUPS=num_groups, BLOCK_G=block_g, NORM_BEFORE_GATE=norm_before_gate, FP8_MIN=fp8_min, FP8_MAX=fp8_max, USE_UE8M0=use_ue8m0, FP8_MIN_SCALING_FACTOR=fp8_min_scaling_factor, num_warps=num_warps, ACTIVATION=activation, ) return x_quant, scales def fused_flatten_fp8_group_quant( x: torch.Tensor, group_size, dtype_quant=fp8_dtype, transpose_scale: bool = False, ): """ Flatten the last two dimension of x and perform FP8 per-token group quantization along the last dimension Key parameters: - x: Matrix X with shape (M, N1, N2). - transpose_scale: If True, return scale with shape (M, cdiv(N1*N2, group_size)) in column-major (transposed) memory layout, i.e. strides (1, M) instead of the default (num_bs_cols, 1). Element values at logical position [m, n] are unchanged; only the physical memory layout differs so downstream consumers (e.g. CK bpreshuffle GEMM) can skip an explicit .transpose(-1, -2).contiguous() before reading. Returns: - out: The output matrix with shape (M, N1 * N2). - out_block_scales: The output matrix with shape (M, cdiv((N1 * N2), group_size)). When transpose_scale=True, strides are (1, M) (column-major); otherwise (num_bs_cols, 1) (row-major). """ M, N1, N2 = x.shape BLOCK_SIZE_N2 = max(triton.next_power_of_2(N2), group_size) N = N1 * N2 num_bs_cols = triton.cdiv(N, group_size) out = torch.empty((M, N), dtype=dtype_quant, device=x.device) if transpose_scale: # Physical buffer is (num_bs_cols, M) row-major; .T gives a # (M, num_bs_cols) view with strides (1, M). The kernel writes # at out_scales_ptr + m * stride_m + n * stride_n, so passing # the natural strides of this view writes to the correct memory # location regardless of layout — no special-case stride wiring # or trailing .view() needed. out_block_scales = torch.empty( (num_bs_cols, M), dtype=torch.float32, device=x.device ).T else: out_block_scales = torch.empty( (M, num_bs_cols), dtype=torch.float32, device=x.device ) DTYPE_MAX = ( torch.finfo(out.dtype).max if torch.is_floating_point(out) else torch.iinfo(out.dtype).max ) grid = ( M, N1, ) _fused_flatten_fp8_group_quant_kernel[grid]( x, out, out_block_scales, *x.stride(), *out.stride(), *out_block_scales.stride(), N2, BLOCK_SIZE_N2=BLOCK_SIZE_N2, QUANT_BLOCK_SIZE=group_size, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, ) return out, out_block_scales def fused_reduce_act_mul_fp8_group_quant( x: torch.Tensor, activation: str = "silu", x2: Optional[torch.Tensor] = None, group_size=128, dtype_quant=fp8_dtype, dtype: Optional[float] = torch.bfloat16, ): """ Apply reduction along the first dimension and apply the activation function + per-token group quantization. If x2 is provided, the only reduction along the first dimension is applied to x2 Args: if x is 3-dim, x: (SPK, M, 2*N1), dtype = fp32. x2: (SPK, M, 2*N1), dtype = fp32. if x is 2-dim, x: (M, 2*N1), dtype = fp16 or bf16. x2 must be None the kernel is essentially identical to aiter.ops.triton.activation.act_mul_and_fp8_group_quant activation: activation function to apply before quantization. - It splits the features into two parts and applies the activation to the first part. - Then, it adds the results together before quantization. - Supports the following activations: - "silu" - "gelu" - "gelu_tanh" Returns: tuple: (y, y_scale), y2 y: (M, N1), dtype = dtype_quant y_scale: (M, cdiv(N1, group_size)), dtype = fp32 y2: (M, N2), dtype = dtype """ _LOGGER.info( f"FUSED_REDUCTION_ACT_MUL_FP8_GROUP_QUANT: x={tuple(x.shape)} activation={activation}" ) assert ( x.dim() == 2 or x.dim() == 3 ), "The number of dimentions for x should be 2 or 3" X_HAS_SPLITK = False x_num_splitk = 1 N2 = 1 y2 = None if x.dim() == 3: x_num_splitk, M, N1 = x.shape x_num_splitk, _, N2 = x2.shape assert ( x.shape[0] == x2.shape[0] and x.shape[1] == x2.shape[1] ), "The first two dimensions should be identical between x and x2" assert ( x_num_splitk > 1 ), "x.shape[0] should be larger then 1 in x.dim() == 3 cases" X_HAS_SPLITK = True y2 = torch.empty((M, N2), dtype=dtype, device=x2.device) else: M, N1 = x.shape assert x2 is None, "x2 should be None in x.dim() == 2 cases" assert ( N1 % 2 == 0 ), "The last dimension for x1 should be multiple of 2 for acitvation and multiplication" N1 = N1 // 2 y = torch.empty((M, N1), dtype=dtype_quant, device=x.device) y_scale = torch.empty( (M, (N1 + group_size - 1) // group_size), dtype=torch.float32, device=x.device, ) BLOCK_SIZE_N1 = max(triton.next_power_of_2(N1), group_size) BLOCK_SIZE_N2 = max(triton.next_power_of_2(N2), 32) BLOCK_SIZE_M2 = 1 if M <= 128 else 4 X_MASK = N1 % BLOCK_SIZE_N1 != 0 DTYPE_MAX = ( torch.finfo(y.dtype).max if torch.is_floating_point(y) else torch.iinfo(y.dtype).max ) num_pid = M if X_HAS_SPLITK: num_pid += triton.cdiv(M, BLOCK_SIZE_M2) * triton.cdiv(N2, BLOCK_SIZE_N2) grid = (num_pid,) _fused_reduce_act_mul_fp8_group_quant[grid]( x, y, y_scale, x2, y2, M, N1, N2, 0 if not X_HAS_SPLITK else x.stride(0), x.stride(0) if not X_HAS_SPLITK else x.stride(1), x.stride(1) if not X_HAS_SPLITK else x.stride(2), y.stride(0), y.stride(1), y_scale.stride(0), y_scale.stride(1), 0 if not X_HAS_SPLITK else x2.stride(0), 0 if not X_HAS_SPLITK else x2.stride(1), 0 if not X_HAS_SPLITK else x2.stride(2), 0 if not X_HAS_SPLITK else y2.stride(0), 0 if not X_HAS_SPLITK else y2.stride(1), ACTIVATION=_get_activation_from_str(activation) if activation else "", BLOCK_SIZE_M2=BLOCK_SIZE_M2, BLOCK_SIZE_N1=BLOCK_SIZE_N1, BLOCK_SIZE_N2=BLOCK_SIZE_N2, QUANT_BLOCK_SIZE=group_size, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, X_HAS_SPLITK=X_HAS_SPLITK, X_NUM_KSPLIT=x_num_splitk, X_NUM_KSPLIT_POW2=triton.next_power_of_2(x_num_splitk), X_MASK=X_MASK, num_warps=1 if max(BLOCK_SIZE_N1, BLOCK_SIZE_N2) <= 512 else 4, ) return (y, y_scale), y2 def fused_reduce_rms_fp8_group_quant( inp1, inp1_weight, inp1_epsilon, inp2=None, inp2_weight=None, inp2_epsilon=None, inp3=None, group_size=128, dtype_quant=fp8_dtype, dtype=None, res1=None, output_unquantized_inp1=False, out3=None, transpose_scale=False, ): """ This op contains several steps: 1. if res1 is not None, inp1 = inp1 + res1, and store inp1 to out_res1 2. perform RMS norm along the last dimenion for inp1 3. if inp2 is not None, perform RMS norm along the last dimenion for inp2 4. perform fp8 quantization for inp1 only 5. if inp3 is not None, perform sum reduction along the first dimension, in the meantime, the inp1 and inp2 has to have the identical first diemsion as inp3 Key parameters: - x: Matrix X with shape (M, N1, N2). Returns: - out1_fp8: The output matrix with shape (M, N1). - out1_bs: The output matrix with shape (M, cdiv(N1, group_size)). - out1: The output matrix with shape (M, N1). - out2: The output matrix with shape (M, N2). - out_res1: The output matrix with shape (M, N1). - out3: The output matrix with shape (M, N3). - out1: The output matrix with shape (M, N1). """ out_dtype = dtype if dtype is not None else inp1.dtype SPK = 1 HAS_SPLITK = False inp1_spk_stride = 0 inp1_row_stride = 0 inp1_col_stride = 0 if inp1.dim() == 3: SPK, M, N1 = inp1.shape assert SPK > 1, "Split-k dimension should have more than 1 element." HAS_SPLITK = True inp1_spk_stride = inp1.stride(0) inp1_row_stride = inp1.stride(1) inp1_col_stride = inp1.stride(2) else: M, N1 = inp1.shape inp1_row_stride = inp1.stride(0) inp1_col_stride = inp1.stride(1) BLOCK_SIZE_N1 = max(triton.next_power_of_2(N1), group_size) if inp2 is not None: if SPK > 1: assert ( inp2.dim() == 3 and inp2.shape[0] == SPK and inp2.shape[1] == M ), f"Incompatible shapes {inp1.shape=}, {inp2.shape=}" _, _, N2 = inp2.shape else: _, N2 = inp2.shape BLOCK_SIZE_N2 = triton.next_power_of_2(N2) else: N2 = 0 BLOCK_SIZE_N2 = 1 if inp3 is not None: assert ( inp3.dim() == 3 and inp3.shape[0] == SPK and inp3.shape[1] == M ), f"Incompatible shapes {inp1.shape=}, {inp3.shape=}" _, _, N3 = inp3.shape BLOCK_SIZE_N3 = triton.next_power_of_2(N3) else: N3 = 0 BLOCK_SIZE_N3 = 1 out1_fp8 = torch.empty((M, N1), dtype=dtype_quant, device=inp1.device) num_bs_cols = (N1 + group_size - 1) // group_size if transpose_scale: # Create with transposed shape for direct transposed storage out1_bs = torch.empty( (num_bs_cols, M), dtype=torch.float32, device=inp1.device, ) else: out1_bs = torch.empty( (M, num_bs_cols), dtype=torch.float32, device=inp1.device, ) out1_fp8_row_stride = out1_fp8.stride(0) out1_fp8_col_stride = out1_fp8.stride(1) # When transpose_scale=True, swap the strides to write directly in transposed layout if transpose_scale: out1_bs_row_stride = out1_bs.stride(1) out1_bs_col_stride = out1_bs.stride(0) else: out1_bs_row_stride = out1_bs.stride(0) out1_bs_col_stride = out1_bs.stride(1) out2 = None inp2_spk_stride = 0 out2_row_stride = 0 out2_col_stride = 0 inp2_row_stride = 0 inp2_col_stride = 0 if inp2 is not None: out2 = torch.empty((M, N2), dtype=out_dtype, device=inp1.device) if SPK > 1: inp2_spk_stride = inp2.stride(0) inp2_row_stride = inp2.stride(1) inp2_col_stride = inp2.stride(2) else: inp2_row_stride = inp2.stride(0) inp2_col_stride = inp2.stride(1) out2_row_stride = out2.stride(0) out2_col_stride = out2.stride(1) inp3_spk_stride = 0 out3_row_stride = 0 out3_col_stride = 0 inp3_row_stride = 0 inp3_col_stride = 0 if inp3 is not None: if out3 is None: out3 = torch.empty((M, N3), dtype=out_dtype, device=inp1.device) inp3_spk_stride = inp3.stride(0) inp3_row_stride = inp3.stride(1) inp3_col_stride = inp3.stride(2) out3_row_stride = out3.stride(0) out3_col_stride = out3.stride(1) out1 = None out1_row_stride = 0 out1_col_stride = 0 if output_unquantized_inp1: out1 = torch.empty((M, N1), dtype=out_dtype, device=inp1.device) out1_row_stride = out1.stride(0) out1_col_stride = out1.stride(1) out_res1 = None res1_row_stride = 0 res1_col_stride = 0 out_res1_row_stride = 0 out_res1_col_stride = 0 if res1 is not None: Mr, Nr = res1.shape assert ( M == Mr and N1 == Nr ), "The shape should be identical between inp1 and res1" out_res1 = torch.empty((M, N1), dtype=out_dtype, device=inp1.device) res1_row_stride = res1.stride(0) res1_col_stride = res1.stride(1) out_res1_row_stride = out_res1.stride(0) out_res1_col_stride = out_res1.stride(1) max_BN = max(BLOCK_SIZE_N1, BLOCK_SIZE_N2, BLOCK_SIZE_N3) if max_BN <= 512: num_warps = 1 elif max_BN <= 2048: num_warps = 4 elif max_BN <= 4096: num_warps = 8 else: num_warps = 16 DTYPE_MAX = ( torch.finfo(out1_fp8.dtype).max if torch.is_floating_point(out1_fp8) else torch.iinfo(out1_fp8.dtype).max ) _fused_reduce_rms_fp8_group_quant_kernel[(3 * M if HAS_SPLITK else 2 * M,)]( inp1, inp1_weight, inp2, inp2_weight, inp3, res1, out1_fp8, out1_bs, out2, out_res1, out1, out3, inp1_epsilon, inp2_epsilon, M, N1, N2, N3, inp1_spk_stride, inp2_spk_stride, inp3_spk_stride, inp1_row_stride, inp2_row_stride, inp3_row_stride, inp1_col_stride, inp2_col_stride, inp3_col_stride, res1_row_stride, res1_col_stride, out1_fp8_row_stride, out1_fp8_col_stride, out1_bs_row_stride, out1_bs_col_stride, out2_row_stride, out2_col_stride, out_res1_row_stride, out_res1_col_stride, out1_row_stride, out1_col_stride, out3_row_stride, out3_col_stride, BLOCK_SIZE_N1=BLOCK_SIZE_N1, BLOCK_SIZE_N2=BLOCK_SIZE_N2, BLOCK_SIZE_N3=BLOCK_SIZE_N3, N_MASK1=(BLOCK_SIZE_N1 != N1), N_MASK2=(BLOCK_SIZE_N2 != N2), N_MASK3=(BLOCK_SIZE_N3 != N3), QUANT_BLOCK_SIZE=group_size, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, HAVE_SECOND_INPUT=(inp2 is not None), FIRST_INPUT_RES=(res1 is not None), FIRST_INPUT_OUT=output_unquantized_inp1, HAS_SPLITK=HAS_SPLITK, NUM_SPLITK=SPK, NUM_SPLITK_POW2=triton.next_power_of_2(SPK), num_warps=num_warps, ) # When transpose_scale=True, view the transposed buffer back to original shape # This keeps shape (M, num_bs_cols) but with column-major memory layout if transpose_scale: out1_bs = out1_bs.view(M, num_bs_cols) return (out1_fp8, out1_bs), out1, out2, out_res1, out3 def fused_silu_mul_fp8_per_tensor_static_quant( inp, inp_scale, dtype_quant=fp8_dtype, silu_convert_to_inp_type=False, ): """ This op contains two steps: 1. compute the silu mul operations 2. perform fp8 quantization for inp1 only Key parameters: - x: Matrix X with shape (M, 2 * N). Returns: - out_fp8: The output matrix with shape (M, N). """ M, N2 = inp.shape assert N2 % 2 == 0 N = N2 // 2 BLOCK_SIZE_N = triton.next_power_of_2(N) out_fp8 = torch.empty((M, N), dtype=dtype_quant, device=inp.device) if BLOCK_SIZE_N <= 512: num_warps = 1 elif BLOCK_SIZE_N <= 2048: num_warps = 4 elif BLOCK_SIZE_N <= 4096: num_warps = 8 else: num_warps = 16 DTYPE_MAX = ( torch.finfo(out_fp8.dtype).max if torch.is_floating_point(out_fp8) else torch.iinfo(out_fp8.dtype).max ) _fused_silu_mul_fp8_per_tensor_static_quant_kernel[(M,)]( inp, out_fp8, inp_scale, M, N, inp.stride(0), inp.stride(1), out_fp8.stride(0), out_fp8.stride(1), BLOCK_SIZE_N=BLOCK_SIZE_N, DTYPE_MAX=DTYPE_MAX, DTYPE_MIN=-DTYPE_MAX, SILU_CONVERT_TO_INP_TYPE=silu_convert_to_inp_type, num_warps=num_warps, ) return out_fp8