Promote latest kernel artifacts to main

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
+*.pyd filter=lfs diff=lfs merge=lfs -text
+*.dylib filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,9 +0,0 @@
-# flashrt/fa2-seqused-runtime
-
-This repository is a compatibility mirror for older `kernels` clients
-that resolve repositories through the default Hugging Face model repo API.
-
-Canonical Kernel Hub repo: https://huggingface.co/kernels/flashrt/fa2-seqused-runtime
-
-Do not edit this mirror by hand. It is generated from the Kernel Hub
-`vN` branches and contains the same `build/**` artifacts.

benchmarks/benchmark.py

@@ -0,0 +1,105 @@
+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 = [
+    # 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),
+]
+
+
+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(
+        "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)
+        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,
+                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__":
+    main()

build/CARD.md

@@ -0,0 +1,58 @@
+---
+tags:
+  - kernel
+  - cuda
+  - flash-attention
+  - inference
+  - cuda-graphs
+library_name: kernels
+---
+
+# FA2 Seqused Runtime
+
+Allocation-free FlashAttention-2 forward operators for CUDA Graph inference.
+The distinguishing feature is a device-resident per-batch `seqused_k`, allowing
+one captured graph to serve changing valid K/V lengths without a host scalar
+read or graph recapture.
+
+## Available functions
+
+- `forward(q, k, v, *, softmax_scale=None, causal=False, use_split_kv=True)`
+- `forward_static(q, k, v, *, out, softmax_lse, workspace=None, softmax_scale=None, causal=False)`
+- `forward_seqused_static(q, k, v, seqused_k, *, out, softmax_lse, workspace=None, softmax_scale=None)`
+- `allocate_outputs(q)`
+- `allocate_workspace(q, k, *, num_sms=None)`
+- `recommended_num_splits(batch, seqlen_q, seqlen_k, heads_q, head_dim, num_sms)`
+- `FA2Workspace`
+
+## Example
+
+```python
+import torch
+from kernels import get_kernel
+
+fa2 = get_kernel("flashrt/fa2-seqused-runtime", version=1)
+q = torch.randn(1, 16, 16, 128, device="cuda", dtype=torch.bfloat16)
+k = torch.randn(1, 2048, 4, 128, device="cuda", dtype=torch.bfloat16)
+v = torch.randn_like(k)
+used = torch.tensor([1536], device="cuda", dtype=torch.int32)
+out, lse = fa2.allocate_outputs(q)
+workspace = fa2.allocate_workspace(q, k)
+
+fa2.forward_seqused_static(
+    q, k, v, used, out=out, softmax_lse=lse, workspace=workspace
+)
+```
+
+The split-KV LSE reset is issued on the current stream and is captured with the
+kernel. Updating `used` on device before replay changes the valid K/V length.
+
+Causal calls use FlashAttention's bottom-right-aligned mask when query and KV
+lengths differ. This is the chunked-prefill/verify convention, not PyTorch
+SDPA's top-left `is_causal=True` convention for rectangular inputs.
+
+## Scope
+
+This package is forward-only and runtime-oriented. It intentionally does not
+duplicate the complete training, varlen, backward, paged-cache, or dropout API
+of `kernels-community/flash-attn2`.

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

@@ -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/torch211-cxx11-cu128-x86_64-linux/_fa2_seqused_runtime_cuda_8fba1d7.abi3.so

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

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

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

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

@@ -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

@@ -0,0 +1,25 @@
+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_8fba1d7",
+  "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_8fba1d7.abi3.so": "qKOG0fqtXr9C5I72edH1AcZsnXm1Rl+rCca/9fwK/Zg=",
+      "_ops.py": "Y+Kfab+oxrJtjqom/l6BdXWyyPHyICYlblsyFF7Ye6Y=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -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/torch211-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_8fba1d7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:566873c30214a6294f174e7d15e1f5240fb491b0c534eb31d154c91900f8f29c
+size 374469512

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

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

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

@@ -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

@@ -0,0 +1,25 @@
+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_8fba1d7",
+  "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_8fba1d7.abi3.so": "VmhzwwIUpilPF059FeH1JA+0kbDFNOsx0VTJGQD48pw=",
+      "_ops.py": "Y+Kfab+oxrJtjqom/l6BdXWyyPHyICYlblsyFF7Ye6Y=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -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/torch212-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_8fba1d7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13c17c724e5c39f2dfa5d6e14b6dae277a207dd65e63af47bd378dee46c8e3bb
+size 374479848

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

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

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

@@ -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

@@ -0,0 +1,25 @@
+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_8fba1d7",
+  "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_8fba1d7.abi3.so": "E8F8ck5cOfLfpdbhS22uJ3ogfdZeY69HvTeN7kbI47s=",
+      "_ops.py": "Y+Kfab+oxrJtjqom/l6BdXWyyPHyICYlblsyFF7Ye6Y=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -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/torch212-cxx11-cu132-x86_64-linux/_fa2_seqused_runtime_cuda_8fba1d7.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38428286c02feeb5039cc0b034f891fc1c938d53ecdfc3d66e67edd483a8afe9
+size 375811104

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

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

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

@@ -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

@@ -0,0 +1,25 @@
+{
+  "name": "fa2-seqused-runtime",
+  "id": "_fa2_seqused_runtime_cuda_8fba1d7",
+  "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_8fba1d7.abi3.so": "OEKChsAv7rUDnMCwNPiR/ByTjVPs38PWbmft1IOor+k=",
+      "_ops.py": "Y+Kfab+oxrJtjqom/l6BdXWyyPHyICYlblsyFF7Ye6Y=",
+      "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
+    }
+  }
+}

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

@@ -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

@@ -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

@@ -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

@@ -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
+    }
+  }
+}

build/torch213-cxx11-cu132-x86_64-linux/__init__.py

@@ -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-cu132-x86_64-linux/_fa2_seqused_runtime_cuda_9ea2146.abi3.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9be77dccfc08ca58a4e17abcaef1b5dbeb8fea968b236b636b6fd17792a9eeeb
+size 375810944

build/torch213-cxx11-cu132-x86_64-linux/_ops.py

@@ -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-cu132-x86_64-linux/metadata.json

@@ -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": "m+d9zPwIylik4Xq8rvG12+uP6paLI2tja2/Rd5Kp7us=",
+      "_ops.py": "csDcHTNxl0h1zDNATeVxZ1/hrMeV5W8mdnPqqupklBk="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81f55ea30fd8f819dcf93a3c934dd584c895bd2f",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "9ea2146c52cd2b967296f47342fb32a14bb25033",
+      "dirty": false
+    }
+  }
+}