liangsu9988 commited on 11 days ago

Commit

8b5c9e3

verified ·

1 Parent(s): e1bbc3c

Uploaded using `kernel-builder`.

Browse files

Files changed (21) hide show

benchmarks/benchmark.py +68 -0
build/torch211-cxx11-cu128-x86_64-linux/__init__.py +284 -0
build/torch211-cxx11-cu128-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so +3 -0
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu128-x86_64-linux/fa2_seqused_runtime/__init__.py +26 -0
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +27 -0
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +284 -0
build/torch211-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so +3 -0
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch211-cxx11-cu130-x86_64-linux/fa2_seqused_runtime/__init__.py +26 -0
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +28 -0
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +284 -0
build/torch212-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so +3 -0
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu130-x86_64-linux/fa2_seqused_runtime/__init__.py +26 -0
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +28 -0
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +284 -0
build/torch212-cxx11-cu132-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so +3 -0
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
build/torch212-cxx11-cu132-x86_64-linux/fa2_seqused_runtime/__init__.py +26 -0
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +28 -0

benchmarks/benchmark.py ADDED Viewed

	@@ -0,0 +1,68 @@

+from __future__ import annotations
+import argparse
+import statistics
+import torch
+import torch.nn.functional as F
+from fa2_seqused_runtime import allocate_outputs, allocate_workspace, forward_static
+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),
+]
+def time_us(fn, warmup=50, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    samples = []
+    for _ in range(repeats):
+        start = torch.cuda.Event(enable_timing=True)
+        end = torch.cuda.Event(enable_timing=True)
+        start.record()
+        fn()
+        end.record()
+        end.synchronize()
+        samples.append(start.elapsed_time(end) * 1000.0)
+    return statistics.median(samples)
+def main():
+    parser = argparse.ArgumentParser()
+    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)
+        out, lse = allocate_outputs(q)
+        workspace = allocate_workspace(q, k)
+        kr = k.repeat_interleave(hq // hkv, dim=2)
+        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__":
+    main()

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

	@@ -0,0 +1,284 @@

+"""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 = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = (96, 128, 256)
+@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(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)
+    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 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,
+    )
+    out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
+        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",
+    "SPLIT_HEAD_DIMS",
+    "SUPPORTED_HEAD_DIMS",
+    "allocate_outputs",
+    "allocate_workspace",
+    "forward",
+    "forward_seqused_static",
+    "forward_static",
+    "recommended_num_splits",
+]

build/torch211-cxx11-cu128-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so ADDED Viewed

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

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

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

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

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

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,27 @@

+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_14f2290",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
+      "_fa2_seqused_runtime_cuda_14f2290.abi3.so": "/po+frEjZ3uIvp1Af1VwvOOQvOd1Kr53biJFsDRAgrw=",
+      "_ops.py": "S1LORAtc4O5qo9aO4HWImXiI6QrfLU+l0PF9N/g+k88=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

	@@ -0,0 +1,284 @@

+"""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 = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = (96, 128, 256)
+@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(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)
+    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 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,
+    )
+    out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
+        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",
+    "SPLIT_HEAD_DIMS",
+    "SUPPORTED_HEAD_DIMS",
+    "allocate_outputs",
+    "allocate_workspace",
+    "forward",
+    "forward_seqused_static",
+    "forward_static",
+    "recommended_num_splits",
+]

build/torch211-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:66afc0c0de8bbe37a508a5ba37e429554b2bfd49a94c417c28aea4cc987cfc38
+size 616325528

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

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

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

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,28 @@

+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_14f2290",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "12.1",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
+      "_fa2_seqused_runtime_cuda_14f2290.abi3.so": "Zq/AwN6LvjelCKW6N+QpVUsr/UmpTEF8KK6kzJh8/Dg=",
+      "_ops.py": "S1LORAtc4O5qo9aO4HWImXiI6QrfLU+l0PF9N/g+k88=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

	@@ -0,0 +1,284 @@

+"""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 = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = (96, 128, 256)
+@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(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)
+    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 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,
+    )
+    out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
+        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",
+    "SPLIT_HEAD_DIMS",
+    "SUPPORTED_HEAD_DIMS",
+    "allocate_outputs",
+    "allocate_workspace",
+    "forward",
+    "forward_seqused_static",
+    "forward_static",
+    "recommended_num_splits",
+]

build/torch212-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so ADDED Viewed

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

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

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

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

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

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

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

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

	@@ -0,0 +1,28 @@

+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_14f2290",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "12.1",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
+      "_fa2_seqused_runtime_cuda_14f2290.abi3.so": "otbTjG0Lkp9V09bUK89p8sv6A0OKM6fngBglzPAn2go=",
+      "_ops.py": "S1LORAtc4O5qo9aO4HWImXiI6QrfLU+l0PF9N/g+k88=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED Viewed

	@@ -0,0 +1,284 @@

+"""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 = (64, 96, 128, 256)
+SPLIT_HEAD_DIMS = (96, 128, 256)
+@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(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)
+    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 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,
+    )
+    out = torch.empty(
+        (splits, q.shape[0], q.shape[2], q.shape[1], q.shape[3]),
+        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",
+    "SPLIT_HEAD_DIMS",
+    "SUPPORTED_HEAD_DIMS",
+    "allocate_outputs",
+    "allocate_workspace",
+    "forward",
+    "forward_seqused_static",
+    "forward_static",
+    "recommended_num_splits",
+]

build/torch212-cxx11-cu132-x86_64-linux/_fa2_seqused_runtime_cuda_14f2290.abi3.so ADDED Viewed

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

+version https://git-lfs.github.com/spec/v1
+oid sha256:81d0f3a0b730aea81498d495c349759895a9076f3d12579837f8a9c636a33282
+size 616122672

build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED Viewed

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

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

build/torch212-cxx11-cu132-x86_64-linux/fa2_seqused_runtime/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED Viewed

	@@ -0,0 +1,28 @@

+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_14f2290",
+  "version": 1,
+  "license": "BSD-3-Clause",
+  "python-depends": [],
+  "backend": {
+    "type": "cuda",
+    "archs": [
+      "10.0",
+      "12.0",
+      "12.1",
+      "8.0",
+      "8.6",
+      "8.9",
+      "9.0"
+    ]
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
+      "_fa2_seqused_runtime_cuda_14f2290.abi3.so": "gdDzoLcwrqgUmNSVw0l1mJWpB289EleYN/ipxjajMoI=",
+      "_ops.py": "S1LORAtc4O5qo9aO4HWImXiI6QrfLU+l0PF9N/g+k88=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}