Uploaded using `kernel-builder`.

Browse files

Files changed (4) hide show

build/torch213-cxx11-cu130-x86_64-linux/__init__.py +293 -0
build/torch213-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_9ea2146.abi3.so +3 -0
build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch213-cxx11-cu130-x86_64-linux/metadata.json +35 -0

build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,293 @@

+"""Static-buffer FlashAttention-2 runtime operators from FlashRT."""
+from __future__ import annotations
+from dataclasses import dataclass
+import math
+from typing import Optional
+import torch
+from ._ops import add_op_namespace_prefix, ops
+SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
+COMPILED_HEAD_DIM_BUCKETS = (64, 96, 128, 256)
+# Vendored FA2 split-KV partial-head tiles are valid only in these logical
+# ranges. Other supported dimensions use the correct no-split kernel.
+SPLIT_HEAD_DIMS = tuple(range(40, 129, 8)) + tuple(range(232, 257, 8))
+@dataclass(frozen=True)
+class FA2Workspace:
+    """Preallocated split-KV workspace for one static attention shape."""
+    softmax_lse_accum: torch.Tensor
+    out_accum: torch.Tensor
+    num_sms: int
+    num_splits: int
+def _ceildiv(a: int, b: int) -> int:
+    return (a + b - 1) // b
+def recommended_num_splits(
+    batch: int,
+    seqlen_q: int,
+    seqlen_k: int,
+    heads_q: int,
+    head_dim: int,
+    num_sms: int,
+) -> int:
+    """Return the exact split count selected by the FlashRT FA2 heuristic."""
+    values = (batch, seqlen_q, seqlen_k, heads_q, head_dim, num_sms)
+    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")
+    if int(head_dim) not in SPLIT_HEAD_DIMS:
+        return 1
+    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)
+    blocks = batch * heads_q * m_blocks
+    effective_sms = num_sms * 2
+    if blocks >= 0.8 * effective_sms:
+        return 1
+    max_splits = min(128, effective_sms, n_blocks)
+    efficiencies = [0.0] * (max_splits + 1)
+    best = 0.0
+    for split in range(1, max_splits + 1):
+        eligible = split == 1 or _ceildiv(n_blocks, split) != _ceildiv(n_blocks, split - 1)
+        if not eligible:
+            continue
+        waves = blocks * split / effective_sms
+        efficiencies[split] = waves / math.ceil(waves)
+        best = max(best, efficiencies[split])
+    for split in range(1, max_splits + 1):
+        eligible = split == 1 or _ceildiv(n_blocks, split) != _ceildiv(n_blocks, split - 1)
+        if eligible and efficiencies[split] >= 0.85 * best:
+            return split
+    return 1
+def allocate_workspace(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    *,
+    num_sms: Optional[int] = None,
+) -> Optional[FA2Workspace]:
+    """Allocate the exact split-KV workspace selected for ``q`` and ``k``.
+    Returns ``None`` when the heuristic selects the no-split path. Allocate
+    once during runtime setup; never call this helper inside a captured loop.
+    """
+    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 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(
+        q.shape[0], q.shape[1], k.shape[1], q.shape[2], q.shape[3], num_sms
+    )
+    if splits == 1:
+        return None
+    lse = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1]),
+        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,
+    )
+    return FA2Workspace(lse, out, int(num_sms), int(splits))
+def allocate_outputs(q: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Allocate output and LSE tensors for a static ``(B,S,H,D)`` query."""
+    if q.ndim != 4:
+        raise ValueError("q must have shape (B, S, H, D)")
+    out = torch.empty_strided(q.shape, q.stride(), device=q.device, dtype=q.dtype)
+    lse = torch.empty(
+        (q.shape[0], q.shape[2], q.shape[1]),
+        device=q.device,
+        dtype=torch.float32,
+    )
+    return out, lse
+def _workspace_args(
+    workspace: Optional[FA2Workspace],
+) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor], int]:
+    if workspace is None:
+        return None, None, 0
+    return workspace.softmax_lse_accum, workspace.out_accum, int(workspace.num_sms)
+@torch.library.register_fake(add_op_namespace_prefix("forward_static"))
+def _forward_static_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    softmax_lse_accum: Optional[torch.Tensor],
+    out_accum: Optional[torch.Tensor],
+    softmax_scale: float,
+    causal: bool = False,
+    num_sms: int = 0,
+) -> None:
+    del k, v, softmax_scale, causal, num_sms
+    if q.ndim != 4 or out.shape != q.shape:
+        raise RuntimeError("q/out must have matching (B, S, H, D) shapes")
+    if softmax_lse.shape != (q.shape[0], q.shape[2], q.shape[1]):
+        raise RuntimeError("softmax_lse must have shape (B, H, S)")
+    if (softmax_lse_accum is None) != (out_accum is None):
+        raise RuntimeError("split-KV workspace tensors must be both set or both None")
+    return None
+@torch.library.register_fake(add_op_namespace_prefix("forward_seqused_static"))
+def _forward_seqused_static_fake(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    seqused_k: torch.Tensor,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    softmax_lse_accum: Optional[torch.Tensor],
+    out_accum: Optional[torch.Tensor],
+    softmax_scale: float,
+    num_sms: int = 0,
+) -> None:
+    del seqused_k
+    return _forward_static_fake(
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        softmax_lse_accum,
+        out_accum,
+        softmax_scale,
+        False,
+        num_sms,
+    )
+def forward_static(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    workspace: Optional[FA2Workspace] = None,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+) -> torch.Tensor:
+    """Run allocation-free FA2 forward into caller-owned static buffers."""
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** -0.5
+    lse_accum, out_accum, num_sms = _workspace_args(workspace)
+    ops.forward_static(
+        q,
+        k,
+        v,
+        out,
+        softmax_lse,
+        lse_accum,
+        out_accum,
+        float(softmax_scale),
+        bool(causal),
+        int(num_sms),
+    )
+    return out
+def forward_seqused_static(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    seqused_k: torch.Tensor,
+    *,
+    out: torch.Tensor,
+    softmax_lse: torch.Tensor,
+    workspace: Optional[FA2Workspace] = None,
+    softmax_scale: Optional[float] = None,
+) -> torch.Tensor:
+    """Run BF16 FA2 with device-resident per-batch K/V lengths.
+    Values in ``seqused_k`` must be in ``[1, k.shape[1]]``. When split-KV is
+    enabled, the LSE workspace is reset to ``-inf`` on the current stream; that
+    reset is captured together with the kernel by CUDA Graphs.
+    """
+    if softmax_scale is None:
+        softmax_scale = q.shape[-1] ** -0.5
+    lse_accum, out_accum, num_sms = _workspace_args(workspace)
+    if lse_accum is not None:
+        lse_accum.fill_(-torch.inf)
+    ops.forward_seqused_static(
+        q,
+        k,
+        v,
+        seqused_k,
+        out,
+        softmax_lse,
+        lse_accum,
+        out_accum,
+        float(softmax_scale),
+        int(num_sms),
+    )
+    return out
+def forward(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    softmax_scale: Optional[float] = None,
+    causal: bool = False,
+    use_split_kv: bool = True,
+) -> torch.Tensor:
+    """Convenience API that allocates outputs and optional split-KV workspace."""
+    out, lse = allocate_outputs(q)
+    workspace = allocate_workspace(q, k) if use_split_kv else None
+    return forward_static(
+        q,
+        k,
+        v,
+        out=out,
+        softmax_lse=lse,
+        workspace=workspace,
+        softmax_scale=softmax_scale,
+        causal=causal,
+    )
+__all__ = [
+    "FA2Workspace",
+    "COMPILED_HEAD_DIM_BUCKETS",
+    "SPLIT_HEAD_DIMS",
+    "SUPPORTED_HEAD_DIMS",
+    "allocate_outputs",
+    "allocate_workspace",
+    "forward",
+    "forward_seqused_static",
+    "forward_static",
+    "recommended_num_splits",
+]

build/torch213-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_9ea2146.abi3.so ADDED Viewed

	@@ -0,0 +1,3 @@

+version https://git-lfs.github.com/spec/v1
+oid sha256:e78b78088c67872a5149377c538bf8c68537469a984f8c34cfd75ae739274d39
+size 374479696

build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED Viewed

	@@ -0,0 +1,9 @@

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

build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,35 @@

+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_9ea2146",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "CN+Ar/JeWCSd1O49S4rR3XPtrtlRXwoY54NuRMsiFDM=",
+      "_fa2_seqused_runtime_cuda_9ea2146.abi3.so": "54t4CIxnhypRSTd8U4v4xoU3RpqYT4w0z9da5zknTTk=",
+      "_ops.py": "csDcHTNxl0h1zDNATeVxZ1/hrMeV5W8mdnPqqupklBk="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "9ea2146c52cd2b967296f47342fb32a14bb25033",
+      "dirty": false
+    }
+  }
+}