Uploaded using `kernel-builder`.

Browse files

Files changed (4) hide show

build/torch211-cxx11-cu128-x86_64-linux/__init__.py +115 -4
build/torch211-cxx11-cu128-x86_64-linux/{_fp8_gemm_cuda_bcb0782.abi3.so → _fp8_gemm_cuda_a6f59d3.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +4 -4

build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED Viewed

@@ -37,6 +37,34 @@ def _fp8_linear_residual_bf16_fake(
     return None
 @torch.library.register_fake(add_op_namespace_prefix("fp8_blockwise_linear_bf16"))
 def _fp8_blockwise_linear_bf16_fake(
     input: torch.Tensor,
@@ -95,9 +123,35 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
     variant = int(variant)
     if m <= 0 or n <= 0 or k <= 0:
         raise RuntimeError("m, n, and k must be positive")
-    if k % 32 != 0:
-        raise RuntimeError("k must be divisible by 32")
     if m == 1:
         if variant == 4:
             return "gemv_fp8_m1_w4"
         if variant == 8:
@@ -113,6 +167,8 @@ def select_fp8_linear_tile(m: int, n: int, k: int, variant: int = 0) -> str:
         return "gemv_fp8_m1_w16"
     if variant != 0:
         raise RuntimeError("small-M dispatcher currently supports variant=0 only")
     if m <= 16:
         if k % 256 == 0:
             return "ld_fp8_gemm_16x128x256_w4" if n % 128 == 0 else "ld_fp8_gemm_16x64x256_w4"
@@ -151,7 +207,9 @@ def fp8_linear_bf16(
     ``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.
     """
     if out is None:
@@ -177,6 +235,56 @@ def fp8_linear_residual_bf16(
     return residual
 def fp8_blockwise_linear_bf16(
     input: torch.Tensor,
     weight: torch.Tensor,
@@ -184,7 +292,7 @@ def fp8_blockwise_linear_bf16(
     weight_scale: torch.Tensor,
     out: torch.Tensor | None = None,
 ) -> torch.Tensor:
-    """Block-128 scaled FP8 linear with BF16 output on SM120."""
     if out is None:
         out = torch.empty(
@@ -228,6 +336,9 @@ def fp8_blockwise_swiglu_quantize_fp8(
 __all__ = [
     "fp8_linear_bf16",
     "fp8_linear_residual_bf16",
     "fp8_blockwise_linear_bf16",
     "fp8_blockwise_swiglu_quantize_fp8",
     "select_fp8_linear_tile",

     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,
     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_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"
     ``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:
     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,
     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(
 __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",

build/torch211-cxx11-cu128-x86_64-linux/{_fp8_gemm_cuda_bcb0782.abi3.so → _fp8_gemm_cuda_a6f59d3.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:182510f121f814d7c367c58c641dc57f8992cc1b62fa25409ed9128e292fcec8
-size 3942856

 version https://git-lfs.github.com/spec/v1
+oid sha256:050b9317e0055fff914a3cf9ff3177bc890b3f29f549780f5feb96e406b58093
+size 3968152

build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _fp8_gemm_cuda_bcb0782
-ops = torch.ops._fp8_gemm_cuda_bcb0782
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_fp8_gemm_cuda_bcb0782::{op_name}"

 import torch
+from . import _fp8_gemm_cuda_a6f59d3
+ops = torch.ops._fp8_gemm_cuda_a6f59d3
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_fp8_gemm_cuda_a6f59d3::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "fp8-gemm",
-  "id": "_fp8_gemm_cuda_bcb0782",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,9 +14,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "AhlkcnyTuk3G4aYsvMqu3QdiEZSTA6Y00poqqiX+Mps=",
-      "_fp8_gemm_cuda_bcb0782.abi3.so": "GCUQ8SH4FNfDZ8WMZB3Ff4mSzBti+iVAntkSjikvzsg=",
-      "_ops.py": "qpBYMOKsfVJtHhKzoX/hxuQPF3MGISv+MQwpPmbzUcw=",
       "fp8_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

 {
   "name": "fp8-gemm",
+  "id": "_fp8_gemm_cuda_a6f59d3",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "HsFfsGY7AK/6Fziv4r2VQ+R2kDjS6hxKKgWqlA+zrKY=",
+      "_fp8_gemm_cuda_a6f59d3.abi3.so": "BQuTF+AFX/+RSjz5/zF3vIkLPyn1SXgPX+uW5Aa1gJM=",
+      "_ops.py": "kgQ6LWi/1cNaF3rEiBHA8t3jyz0toqRVj94wVDPiqVU=",
       "fp8_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }