| """FlashRT FP8 GEMM kernels.""" |
|
|
| from __future__ import annotations |
|
|
| import torch |
|
|
| from ._ops import add_op_namespace_prefix, ops |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bf16")) |
| def _fp8_linear_bf16_fake( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| alpha: float, |
| variant: int, |
| out: torch.Tensor, |
| ) -> None: |
| if input.dim() != 2 or weight.dim() != 2: |
| raise RuntimeError("input and weight must be rank-2 tensors") |
| if out.shape != (input.shape[0], weight.shape[0]): |
| raise RuntimeError("out must have shape (input.shape[0], weight.shape[0])") |
| return None |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_residual_bf16")) |
| def _fp8_linear_residual_bf16_fake( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| alpha: float, |
| variant: int, |
| residual: torch.Tensor, |
| ) -> None: |
| if input.shape[0] != 1: |
| raise RuntimeError("residual path supports only M=1") |
| if residual.shape != (1, weight.shape[0]): |
| raise RuntimeError("residual must have shape (1, weight.shape[0])") |
| return None |
|
|
|
|
| def _check_bias_linear_shapes(input, weight, bias, out) -> None: |
| if input.dim() != 2 or weight.dim() != 2: |
| raise RuntimeError("input and weight must be rank-2 tensors") |
| if input.shape[1] != weight.shape[1]: |
| raise RuntimeError("input and weight K dimensions must match") |
| if bias.shape != (weight.shape[0],): |
| raise RuntimeError("bias must have shape (weight.shape[0],)") |
| if out.shape != (input.shape[0], weight.shape[0]): |
| raise RuntimeError("out must have shape (input.shape[0], weight.shape[0])") |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_bf16")) |
| def _fp8_linear_bias_bf16_fake(input, weight, bias, alpha: float, out) -> None: |
| _check_bias_linear_shapes(input, weight, bias, out) |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_residual_bf16")) |
| def _fp8_linear_bias_residual_bf16_fake( |
| input, weight, bias, alpha: float, residual |
| ) -> None: |
| _check_bias_linear_shapes(input, weight, bias, residual) |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_linear_bias_gelu_bf16")) |
| def _fp8_linear_bias_gelu_bf16_fake(input, weight, bias, alpha: float, out) -> None: |
| _check_bias_linear_shapes(input, weight, bias, out) |
|
|
|
|
| @torch.library.register_fake(add_op_namespace_prefix("fp8_blockwise_linear_bf16")) |
| def _fp8_blockwise_linear_bf16_fake( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| input_scale: torch.Tensor, |
| weight_scale: torch.Tensor, |
| out: torch.Tensor, |
| ) -> None: |
| if input.dim() != 2 or weight.dim() != 2: |
| raise RuntimeError("input and weight must be rank-2 tensors") |
| m, k = input.shape |
| n = weight.shape[0] |
| if weight.shape[1] != k or n % 128 or k % 128: |
| raise RuntimeError("weight shape is invalid or N/K are not divisible by 128") |
| if input_scale.shape != (m, k // 128): |
| raise RuntimeError("input_scale must have shape (M, K / 128)") |
| if weight_scale.shape != (n // 128, k // 128): |
| raise RuntimeError("weight_scale must have shape (N / 128, K / 128)") |
| if out.shape != (m, n): |
| raise RuntimeError("out must have shape (M, N)") |
| return None |
|
|
|
|
| @torch.library.register_fake( |
| add_op_namespace_prefix("fp8_blockwise_swiglu_quantize_fp8") |
| ) |
| def _fp8_blockwise_swiglu_quantize_fp8_fake( |
| input: torch.Tensor, |
| gate_up_weight: torch.Tensor, |
| input_scale: torch.Tensor, |
| gate_up_weight_scale: torch.Tensor, |
| output: torch.Tensor, |
| output_scale: torch.Tensor, |
| ) -> None: |
| m, k = input.shape |
| if gate_up_weight.dim() != 2 or gate_up_weight.shape[0] % 2: |
| raise RuntimeError("gate_up_weight must have shape (2*N, K)") |
| n = gate_up_weight.shape[0] // 2 |
| if gate_up_weight.shape[1] != k or n % 128 or k % 128: |
| raise RuntimeError("gate_up_weight shape is invalid or N/K are not divisible by 128") |
| if input_scale.shape != (m, k // 128): |
| raise RuntimeError("input_scale must have shape (M, K/128)") |
| if gate_up_weight_scale.shape != (2 * n // 128, k // 128): |
| raise RuntimeError("gate_up_weight_scale must have shape (2*N/128, K/128)") |
| if output.shape != (m, n) or output_scale.shape != (m, n // 128): |
| raise RuntimeError("output buffers have invalid shapes") |
| return None |
|
|
|
|
| def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str: |
| """Return the FlashRT tile selected by the public dispatcher.""" |
|
|
| m = int(m) |
| n = int(n) |
| k = int(k) |
| variant = int(variant) |
| if m <= 0 or n <= 0 or k <= 0: |
| raise RuntimeError("m, n, and k must be positive") |
| if k % 16 != 0: |
| raise RuntimeError("k must be divisible by 16") |
| capability = torch.cuda.get_device_capability() if torch.cuda.is_available() else None |
| if capability == (11, 0): |
| forced = {1: "sm110_sq_bf16", 2: "sm110_t1_bf16", 3: "sm110_wide_bf16"} |
| if variant not in {0, *forced}: |
| raise RuntimeError("SM110 variant must be 0 (auto), 1 (Sq), 2 (T1), or 3 (Wide)") |
| if n % 16 or k % 16: |
| raise RuntimeError("SM110 requires n and k divisible by 16") |
| if variant: |
| return forced[variant] |
| if m >= 512 and k == 2048 and 2048 <= n <= 2560: |
| return "sm110_sq_bf16" |
| if m >= 512 and n >= 16 * k: |
| return "sm110_t1_bf16" |
| if m >= 512 and k >= 4 * n: |
| return "sm110_wide_bf16" |
| if n >= 8 * k: |
| return "sm110_wide_bf16" |
| if m >= 128 and k >= 4 * n: |
| return "sm110_sq_bf16" |
| if n == k and m >= 512: |
| return "sm110_sq_bf16" if k <= 1024 else "sm110_wide_bf16" |
| if n == k and m >= 128: |
| return "sm110_wide_bf16" |
| return "sm110_t1_bf16" |
| if m == 1: |
| if k % 32: |
| raise RuntimeError("SM120 requires k divisible by 32") |
| if variant == 4: |
| return "gemv_fp8_m1_w4" |
| if variant == 8: |
| return "gemv_fp8_m1_w8" |
| if variant == 16: |
| return "gemv_fp8_m1_w16" |
| if variant != 0: |
| raise RuntimeError("M=1 variant must be 0, 4, 8, or 16") |
| if n <= 2048: |
| return "gemv_fp8_m1_w4" |
| if n <= 8192: |
| return "gemv_fp8_m1_w8" |
| return "gemv_fp8_m1_w16" |
| if variant != 0: |
| raise RuntimeError("small-M dispatcher currently supports variant=0 only") |
| if k % 32: |
| raise RuntimeError("SM120 requires k divisible by 32") |
| if m <= 16: |
| if k % 256 == 0: |
| return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4" |
| if n % 256 == 0: |
| return "ld_fp8_gemm_16x256x128_w8" |
| if n % 192 == 0: |
| return "ld_fp8_gemm_16x192x128_w4" |
| if n % 128 == 0: |
| return "ld_fp8_gemm_16x128x128_w4" |
| return "ld_fp8_gemm_16x64x128_w4" |
| if m <= 32: |
| if k % 256 == 0: |
| return "ld_fp8_gemm_32x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_32x64x256_w4" |
| if n % 192 == 0: |
| return "ld_fp8_gemm_32x192x128_w4" |
| if n % 128 == 0: |
| return "ld_fp8_gemm_32x128x128_w4" |
| return "ld_fp8_gemm_32x64x128_w4" |
| if m <= 64: |
| if k % 256 == 0: |
| return "ld_fp8_gemm_64x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_64x64x256_w4" |
| if n % 128 == 0: |
| return "ld_fp8_gemm_64x128x128_w4" |
| return "ld_fp8_gemm_64x64x128_w4" |
| raise RuntimeError("only M=1 decode or 2 <= M <= 64 small-M rows are supported") |
|
|
|
|
| def fp8_linear_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| alpha: float = 1.0, |
| out: torch.Tensor | None = None, |
| variant: int = 0, |
| ) -> torch.Tensor: |
| """Compute ``(input @ weight.T) * alpha`` with BF16 output. |
| |
| ``input`` and ``weight`` must be FP8 E4M3 CUDA tensors with shapes |
| ``(M, K)`` and ``(N, K)``. ``alpha`` is a host float, normally the product |
| of static per-tensor input and weight scales. SM110 uses the production |
| CUTLASS Sq/T1/Wide dispatcher over full model row counts; SM120 uses the |
| hand-tuned M<=64 path. |
| """ |
|
|
| if out is None: |
| out = torch.empty( |
| (input.shape[0], weight.shape[0]), |
| device=input.device, |
| dtype=torch.bfloat16, |
| ) |
| ops.fp8_linear_bf16(input, weight, float(alpha), int(variant), out) |
| return out |
|
|
|
|
| def fp8_linear_residual_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| residual: torch.Tensor, |
| alpha: float = 1.0, |
| variant: int = 0, |
| ) -> torch.Tensor: |
| """In-place ``residual += (input @ weight.T) * alpha`` for M=1 decode.""" |
|
|
| ops.fp8_linear_residual_bf16(input, weight, float(alpha), int(variant), residual) |
| return residual |
|
|
|
|
| def fp8_linear_bias_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| bias: torch.Tensor, |
| alpha: float = 1.0, |
| out: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """SM110 FP8 linear with fused BF16 bias and BF16 output.""" |
| if out is None: |
| out = torch.empty( |
| (input.shape[0], weight.shape[0]), |
| device=input.device, |
| dtype=torch.bfloat16, |
| ) |
| ops.fp8_linear_bias_bf16(input, weight, bias, float(alpha), out) |
| return out |
|
|
|
|
| def fp8_linear_bias_residual_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| bias: torch.Tensor, |
| residual: torch.Tensor, |
| alpha: float = 1.0, |
| ) -> torch.Tensor: |
| """SM110 fused ``residual += alpha * input @ weight.T + bias``.""" |
| ops.fp8_linear_bias_residual_bf16( |
| input, weight, bias, float(alpha), residual |
| ) |
| return residual |
|
|
|
|
| def fp8_linear_bias_gelu_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| bias: torch.Tensor, |
| alpha: float = 1.0, |
| out: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """SM110 FP8 linear with fused BF16 bias and GELU epilogue.""" |
| if out is None: |
| out = torch.empty( |
| (input.shape[0], weight.shape[0]), |
| device=input.device, |
| dtype=torch.bfloat16, |
| ) |
| ops.fp8_linear_bias_gelu_bf16(input, weight, bias, float(alpha), out) |
| return out |
|
|
|
|
| def fp8_blockwise_linear_bf16( |
| input: torch.Tensor, |
| weight: torch.Tensor, |
| input_scale: torch.Tensor, |
| weight_scale: torch.Tensor, |
| out: torch.Tensor | None = None, |
| ) -> torch.Tensor: |
| """Block-128 scaled FP8 linear with BF16 output on SM89/SM120.""" |
|
|
| if out is None: |
| out = torch.empty( |
| (input.shape[0], weight.shape[0]), |
| device=input.device, |
| dtype=torch.bfloat16, |
| ) |
| ops.fp8_blockwise_linear_bf16( |
| input, weight, input_scale, weight_scale, out |
| ) |
| return out |
|
|
|
|
| def fp8_blockwise_swiglu_quantize_fp8( |
| input: torch.Tensor, |
| gate_up_weight: torch.Tensor, |
| input_scale: torch.Tensor, |
| gate_up_weight_scale: torch.Tensor, |
| *, |
| output: torch.Tensor | None = None, |
| output_scale: torch.Tensor | None = None, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """SM89 block-128 FP8 gate/up GEMM + SiLU + FP8 requant producer.""" |
|
|
| n = gate_up_weight.shape[0] // 2 |
| if output is None: |
| output = torch.empty( |
| (input.shape[0], n), device=input.device, dtype=torch.float8_e4m3fn |
| ) |
| if output_scale is None: |
| output_scale = torch.empty( |
| (input.shape[0], n // 128), device=input.device, dtype=torch.float32 |
| ) |
| ops.fp8_blockwise_swiglu_quantize_fp8( |
| input, gate_up_weight, input_scale, gate_up_weight_scale, |
| output, output_scale |
| ) |
| return output, output_scale |
|
|
|
|
| __all__ = [ |
| "fp8_linear_bf16", |
| "fp8_linear_residual_bf16", |
| "fp8_linear_bias_bf16", |
| "fp8_linear_bias_residual_bf16", |
| "fp8_linear_bias_gelu_bf16", |
| "fp8_blockwise_linear_bf16", |
| "fp8_blockwise_swiglu_quantize_fp8", |
| "select_fp8_linear_tile", |
| ] |
|
|
