Uploaded using `kernel-builder`.

Browse files

Files changed (5) hide show

benchmarks/benchmark.py +47 -10
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +9 -6
build/torch211-cxx11-cu128-x86_64-linux/{_fa2_seqused_runtime_cuda_61bef7e.abi3.so → _fa2_seqused_runtime_cuda_99d26a1.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

benchmarks/benchmark.py CHANGED Viewed

@@ -10,11 +10,21 @@ from fa2_seqused_runtime import allocate_outputs, allocate_workspace, forward_st
 SHAPES = [
-    (1, 1, 512, 8, 2, 128),
-    (1, 16, 1024, 16, 4, 128),
-    (1, 49, 2520, 24, 4, 128),
-    (1, 64, 4096, 32, 8, 128),
-    (1, 1024, 1024, 32, 8, 128),
 ]
@@ -39,8 +49,13 @@ def main():
     parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16")
     args = parser.parse_args()
     dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
-    print("B,Sq,Sk,Hq,Hkv,D,FlashRT_us,SDPA_expandedGQA_us,Speedup")
-    for batch, sq, sk, hq, hkv, dim in SHAPES:
         q = torch.randn(batch, sq, hq, dim, device="cuda", dtype=dtype)
         k = torch.randn(batch, sk, hkv, dim, device="cuda", dtype=dtype)
         v = torch.randn_like(k)
@@ -50,18 +65,40 @@ def main():
         vr = v.repeat_interleave(hq // hkv, dim=2)
         def flashrt():
-            forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace)
         def sdpa():
-            F.scaled_dot_product_attention(
                 q.permute(0, 2, 1, 3),
                 kr.permute(0, 2, 1, 3),
                 vr.permute(0, 2, 1, 3),
             )
         flashrt_us = time_us(flashrt)
         sdpa_us = time_us(sdpa)
-        print(f"{batch},{sq},{sk},{hq},{hkv},{dim},{flashrt_us:.3f},{sdpa_us:.3f},{sdpa_us / flashrt_us:.3f}")
 if __name__ == "__main__":

 SHAPES = [
+    # GROOT DiT self/cross attention.
+    ("groot-dit-self", 1, 51, 51, 32, 32, 48, False),
+    ("groot-dit-cross", 1, 51, 1024, 32, 32, 48, False),
+    # GROOT N1.7 ViT/VL and SigLIP vision attention.
+    ("groot-n17-vit", 1, 256, 256, 16, 16, 64, False),
+    ("groot-siglip", 2, 256, 256, 16, 16, 72, False),
+    # Qwen2.5-VL/LingBot vision attention.
+    ("vl-vision", 1, 256, 256, 16, 16, 80, False),
+    # Generic runtime/GQA rows.
+    ("gqa-decode", 1, 1, 512, 8, 2, 128, False),
+    ("gqa-short", 1, 16, 1024, 16, 4, 128, False),
+    ("vla-gqa", 1, 49, 2520, 24, 4, 128, False),
+    ("gqa-long-kv", 1, 64, 4096, 32, 8, 128, False),
+    ("qwen-causal", 1, 1024, 1024, 32, 8, 128, True),
+    ("qwen36-causal", 1, 512, 512, 24, 4, 256, True),
 ]
     parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16")
     args = parser.parse_args()
     dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
+    print(
+        "Workload,Mode,B,Sq,Sk,Hq,Hkv,D,FlashRT_us,SDPA_expandedGQA_us,Speedup,"
+        "MaxAbs,P99Abs,MeanAbs,Cosine"
+    )
+    for name, batch, sq, sk, hq, hkv, dim, causal in SHAPES:
+        if causal and dtype == torch.float16:
+            continue
         q = torch.randn(batch, sq, hq, dim, device="cuda", dtype=dtype)
         k = torch.randn(batch, sk, hkv, dim, device="cuda", dtype=dtype)
         v = torch.randn_like(k)
         vr = v.repeat_interleave(hq // hkv, dim=2)
         def flashrt():
+            forward_static(
+                q,
+                k,
+                v,
+                out=out,
+                softmax_lse=lse,
+                workspace=workspace,
+                causal=causal,
+            )
         def sdpa():
+            return F.scaled_dot_product_attention(
                 q.permute(0, 2, 1, 3),
                 kr.permute(0, 2, 1, 3),
                 vr.permute(0, 2, 1, 3),
+                is_causal=causal,
             )
         flashrt_us = time_us(flashrt)
         sdpa_us = time_us(sdpa)
+        actual = out.float()
+        reference = sdpa().permute(0, 2, 1, 3).float()
+        error = (actual - reference).abs()
+        cosine = torch.nn.functional.cosine_similarity(
+            actual.flatten(), reference.flatten(), dim=0
+        ).item()
+        print(
+            f"{name},{'causal' if causal else 'noncausal'},"
+            f"{batch},{sq},{sk},{hq},{hkv},{dim},"
+            f"{flashrt_us:.3f},{sdpa_us:.3f},{sdpa_us / flashrt_us:.3f},"
+            f"{error.max().item():.9f},"
+            f"{torch.quantile(error, 0.99).item():.9f},"
+            f"{error.mean().item():.9f},{cosine:.10f}"
+        )
 if __name__ == "__main__":

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

@@ -11,8 +11,9 @@ import torch
 from ._ops import add_op_namespace_prefix, ops
-SUPPORTED_HEAD_DIMS = (64, 96, 128, 256)
-SPLIT_HEAD_DIMS = (96, 128, 256)
 @dataclass(frozen=True)
@@ -43,7 +44,7 @@ def recommended_num_splits(
     if any(int(v) <= 0 for v in values):
         raise ValueError("all shape values and num_sms must be positive")
     if int(head_dim) not in SUPPORTED_HEAD_DIMS:
-        raise ValueError(f"head_dim must be one of {SUPPORTED_HEAD_DIMS}")
     block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
     n_blocks = _ceildiv(seqlen_k, block_n)
     m_blocks = _ceildiv(seqlen_q, 64)
@@ -82,8 +83,8 @@ def allocate_workspace(
     if q.ndim != 4 or k.ndim != 4:
         raise ValueError("q and k must have shape (B, S, H, D)")
-    if q.shape[-1] not in SPLIT_HEAD_DIMS:
-        return None
     if num_sms is None:
         num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
     splits = recommended_num_splits(
@@ -96,8 +97,9 @@ def allocate_workspace(
         device=q.device,
         dtype=torch.float32,
     )
     out = torch.empty(
-        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
         device=q.device,
         dtype=torch.float32,
     )
@@ -273,6 +275,7 @@ def forward(
 __all__ = [
     "FA2Workspace",
     "SPLIT_HEAD_DIMS",
     "SUPPORTED_HEAD_DIMS",
     "allocate_outputs",

 from ._ops import add_op_namespace_prefix, ops
+SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
+COMPILED_HEAD_DIM_BUCKETS = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = SUPPORTED_HEAD_DIMS
 @dataclass(frozen=True)
     if any(int(v) <= 0 for v in values):
         raise ValueError("all shape values and num_sms must be positive")
     if int(head_dim) not in SUPPORTED_HEAD_DIMS:
+        raise ValueError("head_dim must be a positive multiple of 8 at most 256")
     block_n = 256 if head_dim <= 64 else (128 if head_dim <= 128 else 64)
     n_blocks = _ceildiv(seqlen_k, block_n)
     m_blocks = _ceildiv(seqlen_q, 64)
     if q.ndim != 4 or k.ndim != 4:
         raise ValueError("q and k must have shape (B, S, H, D)")
+    if q.shape[-1] not in SUPPORTED_HEAD_DIMS:
+        raise ValueError("head_dim must be a positive multiple of 8 at most 256")
     if num_sms is None:
         num_sms = torch.cuda.get_device_properties(q.device).multi_processor_count
     splits = recommended_num_splits(
         device=q.device,
         dtype=torch.float32,
     )
+    d_rounded = (q.shape[3] + 31) & ~31
     out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], d_rounded),
         device=q.device,
         dtype=torch.float32,
     )
 __all__ = [
     "FA2Workspace",
+    "COMPILED_HEAD_DIM_BUCKETS",
     "SPLIT_HEAD_DIMS",
     "SUPPORTED_HEAD_DIMS",
     "allocate_outputs",

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

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

 version https://git-lfs.github.com/spec/v1
+oid sha256:824df1266b7f0c7859f9132a97e3f034a9d8b1e434a340386d21a5b533986977
+size 422846752

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

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

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

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

@@ -1,6 +1,6 @@
 {
   "name": "fa2-seqused-runtime",
-  "id": "_fa2_seqused_runtime_cuda_61bef7e",
   "version": 1,
   "license": "BSD-3-Clause",
   "python-depends": [],
@@ -16,9 +16,9 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
-      "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "qmDFPRs3a6HznXIbk9+0J/kRsSDhHWCoVyBdHDRuLFg=",
-      "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
       "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }

 {
   "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_99d26a1",
   "version": 1,
   "license": "BSD-3-Clause",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "zaIQDw3dhB5xbCQPrZ9tGWr/Z559+ooyZOaGIa2/+LQ=",
+      "_fa2_seqused_runtime_cuda_99d26a1.abi3.so": "gk3xJmt/DHhZ+RMql+PwNKnYseQ0o0A4bSGltTOYaXc=",
+      "_ops.py": "3zmyZsIXRYN9PWXNasLSKultcD9S0IaBbtFEFMbUth0=",
       "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
     }
   }