liangsu9988 commited on
Commit
e987d3e
·
verified ·
1 Parent(s): f6d207d

Replace v1 with validated 61bef7e FA2 artifacts

Browse files
build/CARD.md ADDED
@@ -0,0 +1,58 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - kernel
4
+ - cuda
5
+ - flash-attention
6
+ - inference
7
+ - cuda-graphs
8
+ library_name: kernels
9
+ ---
10
+
11
+ # FA2 Seqused Runtime
12
+
13
+ Allocation-free FlashAttention-2 forward operators for CUDA Graph inference.
14
+ The distinguishing feature is a device-resident per-batch `seqused_k`, allowing
15
+ one captured graph to serve changing valid K/V lengths without a host scalar
16
+ read or graph recapture.
17
+
18
+ ## Available functions
19
+
20
+ - `forward(q, k, v, *, softmax_scale=None, causal=False, use_split_kv=True)`
21
+ - `forward_static(q, k, v, *, out, softmax_lse, workspace=None, softmax_scale=None, causal=False)`
22
+ - `forward_seqused_static(q, k, v, seqused_k, *, out, softmax_lse, workspace=None, softmax_scale=None)`
23
+ - `allocate_outputs(q)`
24
+ - `allocate_workspace(q, k, *, num_sms=None)`
25
+ - `recommended_num_splits(batch, seqlen_q, seqlen_k, heads_q, head_dim, num_sms)`
26
+ - `FA2Workspace`
27
+
28
+ ## Example
29
+
30
+ ```python
31
+ import torch
32
+ from kernels import get_kernel
33
+
34
+ fa2 = get_kernel("flashrt/fa2-seqused-runtime", version=1)
35
+ q = torch.randn(1, 16, 16, 128, device="cuda", dtype=torch.bfloat16)
36
+ k = torch.randn(1, 2048, 4, 128, device="cuda", dtype=torch.bfloat16)
37
+ v = torch.randn_like(k)
38
+ used = torch.tensor([1536], device="cuda", dtype=torch.int32)
39
+ out, lse = fa2.allocate_outputs(q)
40
+ workspace = fa2.allocate_workspace(q, k)
41
+
42
+ fa2.forward_seqused_static(
43
+ q, k, v, used, out=out, softmax_lse=lse, workspace=workspace
44
+ )
45
+ ```
46
+
47
+ The split-KV LSE reset is issued on the current stream and is captured with the
48
+ kernel. Updating `used` on device before replay changes the valid K/V length.
49
+
50
+ Causal calls use FlashAttention's bottom-right-aligned mask when query and KV
51
+ lengths differ. This is the chunked-prefill/verify convention, not PyTorch
52
+ SDPA's top-left `is_causal=True` convention for rectangular inputs.
53
+
54
+ ## Scope
55
+
56
+ This package is forward-only and runtime-oriented. It intentionally does not
57
+ duplicate the complete training, varlen, backward, paged-cache, or dropout API
58
+ of `kernels-community/flash-attn2`.
build/benchmarks/README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Benchmarks
2
+
3
+ ```bash
4
+ python benchmarks/benchmark.py --dtype bf16
5
+ python benchmarks/benchmark.py --dtype fp16
6
+ ```
7
+
8
+ The FlashRT timing includes split-KV dispatch when the package heuristic selects
9
+ it. The SDPA baseline explicitly materializes repeated K/V heads for GQA, so
10
+ results must label that fact. Installed-artifact results are recorded only after
11
+ the corresponding Hub build is available.
build/benchmarks/RESULTS.md ADDED
@@ -0,0 +1,35 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Benchmark results
2
+
3
+ ## Pre-release source qualification
4
+
5
+ - GPU: NVIDIA GeForce RTX 5090, SM120
6
+ - PyTorch: `2.9.0a0+145a3a7bda.nv25.10`
7
+ - CUDA: 13.0
8
+ - FlashRT source: `b3eab55`
9
+ - Dtype: BF16
10
+ - Timing: warmed CUDA events, median of 20 groups x 20 calls
11
+ - Baseline: PyTorch SDPA with GQA heads expanded before timing
12
+
13
+ | B | Sq | Sk | Hq/Hkv | D | Splits | Package us | Original FlashRT us | SDPA us | vs original | vs SDPA |
14
+ |---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|---:|
15
+ | 1 | 1 | 512 | 8/2 | 128 | 4 | 10.453 | 11.340 | 11.684 | 0.922x | 1.118x |
16
+ | 1 | 16 | 1024 | 16/4 | 128 | 8 | 10.477 | 11.298 | 11.712 | 0.927x | 1.118x |
17
+ | 1 | 49 | 2520 | 24/4 | 128 | 10 | 22.761 | 23.014 | 25.011 | 0.989x | 1.099x |
18
+ | 1 | 64 | 4096 | 32/8 | 128 | 7 | 43.226 | 43.461 | 43.662 | 0.995x | 1.010x |
19
+ | 1 | 1024 | 1024 | 32/8 | 128 | 1 | 108.701 | 108.928 | 111.005 | 0.998x | 1.021x |
20
+
21
+ `vs original` is `package/original`; values at or below 1.0 show no packaged
22
+ kernel regression. `vs SDPA` is `SDPA/package`. K/V head expansion is excluded
23
+ from SDPA timing, so the comparison does not credit FlashRT for avoiding that
24
+ materialization.
25
+
26
+ The 40-row detailed accuracy sweep over FP16/BF16, D=64/96/128/256, MHA/GQA,
27
+ partial tiles and causal BF16 produced:
28
+
29
+ - worst maximum absolute error vs PyTorch SDPA: `0.001953125`
30
+ - worst cosine similarity: `0.9999961853`
31
+ - package output vs original FlashRT output on benchmark rows: exact
32
+
33
+ These are source-artifact qualification numbers, not the final Hub artifact
34
+ claim. Installed `kernel-builder` results replace this section after upload and
35
+ fresh-process `get_kernel(..., version=1)` validation.
build/benchmarks/benchmark.py ADDED
@@ -0,0 +1,68 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from __future__ import annotations
2
+
3
+ import argparse
4
+ import statistics
5
+
6
+ import torch
7
+ import torch.nn.functional as F
8
+
9
+ from fa2_seqused_runtime import allocate_outputs, allocate_workspace, forward_static
10
+
11
+
12
+ SHAPES = [
13
+ (1, 1, 512, 8, 2, 128),
14
+ (1, 16, 1024, 16, 4, 128),
15
+ (1, 49, 2520, 24, 4, 128),
16
+ (1, 64, 4096, 32, 8, 128),
17
+ (1, 1024, 1024, 32, 8, 128),
18
+ ]
19
+
20
+
21
+ def time_us(fn, warmup=50, repeats=200):
22
+ for _ in range(warmup):
23
+ fn()
24
+ torch.cuda.synchronize()
25
+ samples = []
26
+ for _ in range(repeats):
27
+ start = torch.cuda.Event(enable_timing=True)
28
+ end = torch.cuda.Event(enable_timing=True)
29
+ start.record()
30
+ fn()
31
+ end.record()
32
+ end.synchronize()
33
+ samples.append(start.elapsed_time(end) * 1000.0)
34
+ return statistics.median(samples)
35
+
36
+
37
+ def main():
38
+ parser = argparse.ArgumentParser()
39
+ parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16")
40
+ args = parser.parse_args()
41
+ dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
42
+ print("B,Sq,Sk,Hq,Hkv,D,FlashRT_us,SDPA_expandedGQA_us,Speedup")
43
+ for batch, sq, sk, hq, hkv, dim in SHAPES:
44
+ q = torch.randn(batch, sq, hq, dim, device="cuda", dtype=dtype)
45
+ k = torch.randn(batch, sk, hkv, dim, device="cuda", dtype=dtype)
46
+ v = torch.randn_like(k)
47
+ out, lse = allocate_outputs(q)
48
+ workspace = allocate_workspace(q, k)
49
+ kr = k.repeat_interleave(hq // hkv, dim=2)
50
+ vr = v.repeat_interleave(hq // hkv, dim=2)
51
+
52
+ def flashrt():
53
+ forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace)
54
+
55
+ def sdpa():
56
+ F.scaled_dot_product_attention(
57
+ q.permute(0, 2, 1, 3),
58
+ kr.permute(0, 2, 1, 3),
59
+ vr.permute(0, 2, 1, 3),
60
+ )
61
+
62
+ flashrt_us = time_us(flashrt)
63
+ sdpa_us = time_us(sdpa)
64
+ print(f"{batch},{sq},{sk},{hq},{hkv},{dim},{flashrt_us:.3f},{sdpa_us:.3f},{sdpa_us / flashrt_us:.3f}")
65
+
66
+
67
+ if __name__ == "__main__":
68
+ main()
build/torch211-cxx11-cu128-x86_64-linux/_fa2_seqused_runtime_cuda_61bef7e.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aa60c53d1b376ba1f39d721b93dfb427f911b120e11d60a857205d1c346e2c58
3
+ size 391476840
build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fa2_seqused_runtime_cuda_e9a7efd
3
- ops = torch.ops._fa2_seqused_runtime_cuda_e9a7efd
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fa2_seqused_runtime_cuda_e9a7efd::{op_name}"
 
1
  import torch
2
+ from . import _fa2_seqused_runtime_cuda_61bef7e
3
+ ops = torch.ops._fa2_seqused_runtime_cuda_61bef7e
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fa2_seqused_runtime_cuda_61bef7e::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fa2-seqused-runtime",
3
- "id": "_fa2_seqused_runtime_cuda_e9a7efd",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
@@ -10,8 +10,6 @@
10
  "10.0",
11
  "12.0",
12
  "8.0",
13
- "8.6",
14
- "8.9",
15
  "9.0"
16
  ]
17
  },
@@ -19,8 +17,8 @@
19
  "algorithm": "sha256",
20
  "files": {
21
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
22
- "_fa2_seqused_runtime_cuda_e9a7efd.abi3.so": "+hkDrx81EIUt9ECTogeAn4XyulcbdPMG144yfnblbaI=",
23
- "_ops.py": "LUT5Vb7hzmHsRm5518t44oMDhSUB+6QoYZn/pY5QEas=",
24
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
25
  }
26
  }
 
1
  {
2
  "name": "fa2-seqused-runtime",
3
+ "id": "_fa2_seqused_runtime_cuda_61bef7e",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
 
10
  "10.0",
11
  "12.0",
12
  "8.0",
 
 
13
  "9.0"
14
  ]
15
  },
 
17
  "algorithm": "sha256",
18
  "files": {
19
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
20
+ "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "qmDFPRs3a6HznXIbk9+0J/kRsSDhHWCoVyBdHDRuLFg=",
21
+ "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }
build/torch211-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_61bef7e.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:414ce4814305eb711f6b58d804ddb0539e0b6e70685679070369c94938324968
3
+ size 358207912
build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fa2_seqused_runtime_cuda_e9a7efd
3
- ops = torch.ops._fa2_seqused_runtime_cuda_e9a7efd
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fa2_seqused_runtime_cuda_e9a7efd::{op_name}"
 
1
  import torch
2
+ from . import _fa2_seqused_runtime_cuda_61bef7e
3
+ ops = torch.ops._fa2_seqused_runtime_cuda_61bef7e
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fa2_seqused_runtime_cuda_61bef7e::{op_name}"
build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fa2-seqused-runtime",
3
- "id": "_fa2_seqused_runtime_cuda_e9a7efd",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
@@ -9,10 +9,7 @@
9
  "archs": [
10
  "10.0",
11
  "12.0",
12
- "12.1",
13
  "8.0",
14
- "8.6",
15
- "8.9",
16
  "9.0"
17
  ]
18
  },
@@ -20,8 +17,8 @@
20
  "algorithm": "sha256",
21
  "files": {
22
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
23
- "_fa2_seqused_runtime_cuda_e9a7efd.abi3.so": "/6dGG6x2WpMhUJfwZYI+ZZb9lNbFDI8k4RX3O21Hzb4=",
24
- "_ops.py": "LUT5Vb7hzmHsRm5518t44oMDhSUB+6QoYZn/pY5QEas=",
25
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
26
  }
27
  }
 
1
  {
2
  "name": "fa2-seqused-runtime",
3
+ "id": "_fa2_seqused_runtime_cuda_61bef7e",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
 
9
  "archs": [
10
  "10.0",
11
  "12.0",
 
12
  "8.0",
 
 
13
  "9.0"
14
  ]
15
  },
 
17
  "algorithm": "sha256",
18
  "files": {
19
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
20
+ "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "QUzkgUMF63Efa1jYBN2wU54LbnBoVnkHA2nJSTgySWg=",
21
+ "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }
build/torch212-cxx11-cu130-x86_64-linux/_fa2_seqused_runtime_cuda_61bef7e.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:31ade01e2d39da172b84c435c935055ffed2cf25ecf9378ee46e6a3129350a1f
3
+ size 358213904
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fa2_seqused_runtime_cuda_e9a7efd
3
- ops = torch.ops._fa2_seqused_runtime_cuda_e9a7efd
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fa2_seqused_runtime_cuda_e9a7efd::{op_name}"
 
1
  import torch
2
+ from . import _fa2_seqused_runtime_cuda_61bef7e
3
+ ops = torch.ops._fa2_seqused_runtime_cuda_61bef7e
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fa2_seqused_runtime_cuda_61bef7e::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fa2-seqused-runtime",
3
- "id": "_fa2_seqused_runtime_cuda_e9a7efd",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
@@ -9,10 +9,7 @@
9
  "archs": [
10
  "10.0",
11
  "12.0",
12
- "12.1",
13
  "8.0",
14
- "8.6",
15
- "8.9",
16
  "9.0"
17
  ]
18
  },
@@ -20,8 +17,8 @@
20
  "algorithm": "sha256",
21
  "files": {
22
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
23
- "_fa2_seqused_runtime_cuda_e9a7efd.abi3.so": "S2oc2AfTD43/13DedBbf58FkbLuKlCYTb7tCCQvLN08=",
24
- "_ops.py": "LUT5Vb7hzmHsRm5518t44oMDhSUB+6QoYZn/pY5QEas=",
25
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
26
  }
27
  }
 
1
  {
2
  "name": "fa2-seqused-runtime",
3
+ "id": "_fa2_seqused_runtime_cuda_61bef7e",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
 
9
  "archs": [
10
  "10.0",
11
  "12.0",
 
12
  "8.0",
 
 
13
  "9.0"
14
  ]
15
  },
 
17
  "algorithm": "sha256",
18
  "files": {
19
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
20
+ "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "Ma3gHi052hcrhMQ1yTUFX/7SzyXs+TeO5G5qMSk1Ch8=",
21
+ "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }
build/torch212-cxx11-cu132-x86_64-linux/_fa2_seqused_runtime_cuda_61bef7e.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:3fe62ce8f38b95cce80c8b86b64a52d34c16d6c588693c7211f3cfe56279148a
3
+ size 358869320
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fa2_seqused_runtime_cuda_e9a7efd
3
- ops = torch.ops._fa2_seqused_runtime_cuda_e9a7efd
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fa2_seqused_runtime_cuda_e9a7efd::{op_name}"
 
1
  import torch
2
+ from . import _fa2_seqused_runtime_cuda_61bef7e
3
+ ops = torch.ops._fa2_seqused_runtime_cuda_61bef7e
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fa2_seqused_runtime_cuda_61bef7e::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fa2-seqused-runtime",
3
- "id": "_fa2_seqused_runtime_cuda_e9a7efd",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
@@ -9,10 +9,7 @@
9
  "archs": [
10
  "10.0",
11
  "12.0",
12
- "12.1",
13
  "8.0",
14
- "8.6",
15
- "8.9",
16
  "9.0"
17
  ]
18
  },
@@ -20,8 +17,8 @@
20
  "algorithm": "sha256",
21
  "files": {
22
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
23
- "_fa2_seqused_runtime_cuda_e9a7efd.abi3.so": "DdVDYHNpVJs/yyigwRVVaxid2dJt+gNf019kM/4aWTI=",
24
- "_ops.py": "LUT5Vb7hzmHsRm5518t44oMDhSUB+6QoYZn/pY5QEas=",
25
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
26
  }
27
  }
 
1
  {
2
  "name": "fa2-seqused-runtime",
3
+ "id": "_fa2_seqused_runtime_cuda_61bef7e",
4
  "version": 1,
5
  "license": "BSD-3-Clause",
6
  "python-depends": [],
 
9
  "archs": [
10
  "10.0",
11
  "12.0",
 
12
  "8.0",
 
 
13
  "9.0"
14
  ]
15
  },
 
17
  "algorithm": "sha256",
18
  "files": {
19
  "__init__.py": "ck4+/aHijqRRXtVYRzkkIj7bHwvMxxOEnRIKVQSkd6k=",
20
+ "_fa2_seqused_runtime_cuda_61bef7e.abi3.so": "P+Ys6POLlczoDIuGtkpS00wW1sWIaTxyEfPP5WJ5FIo=",
21
+ "_ops.py": "S5eXSORqSBWo6D7FltYgk8ttuDlbt6lSeHd0GmU8q6U=",
22
  "fa2_seqused_runtime/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
23
  }
24
  }