fold-attention 0.1.0__tar.gz

fold_attention-0.1.0/PKG-INFO

@@ -0,0 +1,124 @@
+Metadata-Version: 2.3
+Name: fold-attention
+Version: 0.1.0
+Summary: FoldAttention declares softmax’s reference before execution, making attention additive and enabling up to 3.09× faster H100 decode and deterministic backward faster than nondeterministic kernels.
+Author: Sriman Achanta
+Author-email: Sriman Achanta <srimanachanta@gmail.com>
+Requires-Dist: flash-attn-4>=4.0.0b31
+Requires-Dist: numpy>=2.5.3
+Requires-Dist: nvidia-cutlass-dsl>=4.7.1
+Requires-Dist: quack-kernels>=0.6.5
+Requires-Dist: torch>=2.14.0
+Requires-Dist: flash-attn-4[cu13]>=4.0.0b31 ; extra == 'cu13'
+Requires-Dist: nvidia-cutlass-dsl[cu13]>=4.7.1 ; extra == 'cu13'
+Requires-Dist: quack-kernels[cu13]>=0.6.5 ; extra == 'cu13'
+Requires-Python: >=3.12
+Provides-Extra: cu13
+Description-Content-Type: text/markdown
+
+# FoldAttention
+
+<p align="center">
+  <img src="assets/hero.png" alt="FoldAttention overview: declared-reference softmax, decode speedup, and deterministic backward throughput">
+</p>
+
+FoldAttention declares softmax's reference before execution. Final weights
+make decode contributions additive, while a shared integer grid makes the
+backward deterministic without serializing its reductions.
+
+## Usage
+
+Install FoldAttention from PyPI:
+
+```bash
+pip install fold-attention
+```
+
+Install the CUDA 13 dependencies with the `cu13` extra:
+
+```bash
+pip install "fold-attention[cu13]"
+```
+
+FoldAttention requires Python 3.12 or newer and an NVIDIA SM90 GPU.
+
+## Code usage
+
+### Training
+
+`fold_attn_func` follows FlashAttention's `(batch, seqlen, heads, head_dim)`
+layout. Keys and values may use fewer heads for GQA or MQA.
+
+```python
+from fold_attention import fold_attn_func
+
+out = fold_attn_func(q, k, v, causal=True)
+out.backward(dout)
+```
+
+The forward uses FlashAttention-4. The FoldAttention backward produces
+bit-identical gradients across repeated runs, batching, and variable-length
+packing.
+
+For packed self-attention, pass `(total_tokens, heads, head_dim)` tensors and
+a CUDA `int32` cumulative-length vector:
+
+```python
+from fold_attention import fold_attn_varlen_func
+
+out = fold_attn_varlen_func(q, k, v, cu_seqlens, causal=True)
+```
+
+### Decode
+
+`FoldKVCache` owns one layer's paged cache. `prefill` runs FlashAttention-4
+and writes the cache. `fold_attn_with_kvcache` optionally appends one token per
+request, then attends over the updated cache.
+
+```python
+from fold_attention import FoldKVCache, fold_attn_with_kvcache
+
+cache = FoldKVCache(
+    batch=batch_size,
+    n_heads=n_heads,
+    n_kv_heads=n_kv_heads,
+    head_dim=head_dim,
+    max_len=max_len,
+    depth=16,
+)
+
+prompt_out = cache.prefill(q, k, v, cu_seqlens)
+step_out = fold_attn_with_kvcache(q_step, cache, k_step, v_step)
+```
+
+Set `depth=None` for dense decode. A finite depth cuts low-weight keys while
+retaining their normalization mass. Set `v8=True` to store values in two E4M3
+planes.
+
+## Benchmarks
+
+### Decode speed and accuracy
+
+Latency against FP32-relative error on seven real-model generations. Each
+FoldAttention curve sweeps the decode depth.
+
+<p align="center">
+  <img src="assets/decode-pareto.png" alt="Decode latency against FP32-relative error on seven real-model generations">
+</p>
+
+### Deterministic backward
+
+Causal backward throughput across MHA and GQA shapes. The rows below each
+panel report FoldAttention throughput relative to the fastest deterministic
+and nondeterministic kernel.
+
+<p align="center">
+  <img src="assets/backward-throughput.png" alt="Causal attention backward throughput on MHA and GQA shapes">
+</p>
+
+The benchmark suite and measurement protocol are documented in
+[`benchmarks/README.md`](benchmarks/README.md).
+
+## License
+
+FoldAttention is released under the [Apache License 2.0](LICENSE).

fold_attention-0.1.0/README.md

@@ -0,0 +1,106 @@
+# FoldAttention
+
+<p align="center">
+  <img src="assets/hero.png" alt="FoldAttention overview: declared-reference softmax, decode speedup, and deterministic backward throughput">
+</p>
+
+FoldAttention declares softmax's reference before execution. Final weights
+make decode contributions additive, while a shared integer grid makes the
+backward deterministic without serializing its reductions.
+
+## Usage
+
+Install FoldAttention from PyPI:
+
+```bash
+pip install fold-attention
+```
+
+Install the CUDA 13 dependencies with the `cu13` extra:
+
+```bash
+pip install "fold-attention[cu13]"
+```
+
+FoldAttention requires Python 3.12 or newer and an NVIDIA SM90 GPU.
+
+## Code usage
+
+### Training
+
+`fold_attn_func` follows FlashAttention's `(batch, seqlen, heads, head_dim)`
+layout. Keys and values may use fewer heads for GQA or MQA.
+
+```python
+from fold_attention import fold_attn_func
+
+out = fold_attn_func(q, k, v, causal=True)
+out.backward(dout)
+```
+
+The forward uses FlashAttention-4. The FoldAttention backward produces
+bit-identical gradients across repeated runs, batching, and variable-length
+packing.
+
+For packed self-attention, pass `(total_tokens, heads, head_dim)` tensors and
+a CUDA `int32` cumulative-length vector:
+
+```python
+from fold_attention import fold_attn_varlen_func
+
+out = fold_attn_varlen_func(q, k, v, cu_seqlens, causal=True)
+```
+
+### Decode
+
+`FoldKVCache` owns one layer's paged cache. `prefill` runs FlashAttention-4
+and writes the cache. `fold_attn_with_kvcache` optionally appends one token per
+request, then attends over the updated cache.
+
+```python
+from fold_attention import FoldKVCache, fold_attn_with_kvcache
+
+cache = FoldKVCache(
+    batch=batch_size,
+    n_heads=n_heads,
+    n_kv_heads=n_kv_heads,
+    head_dim=head_dim,
+    max_len=max_len,
+    depth=16,
+)
+
+prompt_out = cache.prefill(q, k, v, cu_seqlens)
+step_out = fold_attn_with_kvcache(q_step, cache, k_step, v_step)
+```
+
+Set `depth=None` for dense decode. A finite depth cuts low-weight keys while
+retaining their normalization mass. Set `v8=True` to store values in two E4M3
+planes.
+
+## Benchmarks
+
+### Decode speed and accuracy
+
+Latency against FP32-relative error on seven real-model generations. Each
+FoldAttention curve sweeps the decode depth.
+
+<p align="center">
+  <img src="assets/decode-pareto.png" alt="Decode latency against FP32-relative error on seven real-model generations">
+</p>
+
+### Deterministic backward
+
+Causal backward throughput across MHA and GQA shapes. The rows below each
+panel report FoldAttention throughput relative to the fastest deterministic
+and nondeterministic kernel.
+
+<p align="center">
+  <img src="assets/backward-throughput.png" alt="Causal attention backward throughput on MHA and GQA shapes">
+</p>
+
+The benchmark suite and measurement protocol are documented in
+[`benchmarks/README.md`](benchmarks/README.md).
+
+## License
+
+FoldAttention is released under the [Apache License 2.0](LICENSE).

fold_attention-0.1.0/pyproject.toml

@@ -0,0 +1,62 @@
+[project]
+name = "fold-attention"
+version = "0.1.0"
+description = "FoldAttention declares softmax’s reference before execution, making attention additive and enabling up to 3.09× faster H100 decode and deterministic backward faster than nondeterministic kernels."
+readme = "README.md"
+requires-python = ">=3.12"
+dependencies = [
+    "flash-attn-4>=4.0.0b31",
+    "numpy>=2.5.3",
+    "nvidia-cutlass-dsl>=4.7.1",
+    "quack-kernels>=0.6.5",
+    "torch>=2.14.0",
+]
+
+[[project.authors]]
+name = "Sriman Achanta"
+email = "srimanachanta@gmail.com"
+
+[project.optional-dependencies]
+cu13 = [
+    "flash-attn-4[cu13]>=4.0.0b31",
+    "nvidia-cutlass-dsl[cu13]>=4.7.1",
+    "quack-kernels[cu13]>=0.6.5",
+]
+
+[dependency-groups]
+dev = ["pytest>=8.0"]
+
+[build-system]
+requires = ["uv_build>=0.12.10,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.sources.torch]
+index = "pytorch-cu130"
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.ty.environment]
+python = ".venv"
+extra-paths = [".venv/lib/python3.14/site-packages/nvidia_cutlass_dsl/dsl_packages"]
+
+[tool.ty.analysis]
+allowed-unresolved-imports = ["cuda.bindings.**"]
+replace-imports-with-any = ["cutlass._mlir.**"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py312"
+
+[tool.ruff.lint]
+ignore = [
+    "C408",
+    "FURB136",
+    "PLR1730",
+    "SIM102",
+    "SIM108",
+    "SIM113",
+    "SIM114",
+]

fold_attention-0.1.0/pyproject.toml.orig

@@ -0,0 +1,74 @@
+[project]
+name = "fold-attention"
+version = "0.1.0"
+description = "FoldAttention declares softmax’s reference before execution, making attention additive and enabling up to 3.09× faster H100 decode and deterministic backward faster than nondeterministic kernels."
+readme = "README.md"
+authors = [
+    { name = "Sriman Achanta", email = "srimanachanta@gmail.com" }
+]
+requires-python = ">=3.12"
+dependencies = [
+    "flash-attn-4>=4.0.0b31",
+    "numpy>=2.5.3",
+    "nvidia-cutlass-dsl>=4.7.1",
+    "quack-kernels>=0.6.5",
+    "torch>=2.14.0",
+]
+
+[dependency-groups]
+dev = [
+    "pytest>=8.0",
+]
+
+[project.optional-dependencies]
+cu13 = [
+    "flash-attn-4[cu13]>=4.0.0b31",
+    "nvidia-cutlass-dsl[cu13]>=4.7.1",
+    "quack-kernels[cu13]>=0.6.5",
+]
+
+[build-system]
+requires = ["uv_build>=0.12.10,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.sources]
+torch = { index = "pytorch-cu130" }
+
+[[tool.uv.index]]
+name = "pytorch-cu130"
+url = "https://download.pytorch.org/whl/cu130"
+explicit = true
+
+[tool.ty.environment]
+python = ".venv"
+extra-paths = [
+    ".venv/lib/python3.14/site-packages/nvidia_cutlass_dsl/dsl_packages"
+]
+
+[tool.ty.analysis]
+allowed-unresolved-imports = [
+    "cuda.bindings.**",
+]
+replace-imports-with-any = ["cutlass._mlir.**"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py312"
+
+[tool.ruff.lint]
+ignore = [
+    # keyword dicts read as the calls they feed
+    "C408",
+    # Rewrites that break or change CuTe DSL tracing. A condition on a traced
+    # value has no Python bool, so builtin `min`/`max` (PLR1730, FURB136) and
+    # a ternary (SIM108) fail to compile; `enumerate` hides a `cutlass.range`
+    # loop from the DSL's loop rewrite (SIM113); collapsing or merging `if`
+    # statements (SIM102, SIM114) changes the regions a dynamic condition
+    # traces into.
+    "FURB136",
+    "PLR1730",
+    "SIM102",
+    "SIM108",
+    "SIM113",
+    "SIM114",
+]

fold_attention-0.1.0/src/fold_attention/__init__.py

@@ -0,0 +1,16 @@
+"""FoldAttention: softmax attention whose cross-CTA reductions do not depend
+on the order work runs in.
+
+Training uses FlashAttention-4's forward and an SM90 backward whose dQ, dK and
+dV sums are integer or fixed-order reductions, so gradients are the same bits
+on every run and for every batching of a request. Serving decodes over a
+quantised paged cache against a static per-row reference, so split-KV,
+shared-prefix and speculative partials combine by plain addition.
+"""
+
+from .interface import fold_attn_func, fold_attn_varlen_func, fold_attn_with_kvcache
+from .kv_cache import FoldKVCache
+
+__version__ = "0.1.0"
+
+__all__ = ["FoldKVCache", "fold_attn_func", "fold_attn_varlen_func", "fold_attn_with_kvcache"]

fold_attention-0.1.0/src/fold_attention/backward/__init__.py

@@ -0,0 +1,6 @@
+"""The attention backward with order-free cross-CTA reductions."""
+
+from .launch import prepare_backward
+from .plan import Plan, plan_dense, plan_varlen, tile_config
+
+__all__ = ["Plan", "plan_dense", "plan_varlen", "prepare_backward", "tile_config"]

fold_attention-0.1.0/src/fold_attention/backward/ablation.py

@@ -0,0 +1,29 @@
+"""The backward with one mechanism changed, for measuring what it buys.
+
+Each variant is the shipped kernel (`prepare_backward`) with one choice
+replaced and the rest as it is:
+
+- `fp32_dq=True`: dQ's partials summed as fp32 by `cp.reduce.async.bulk
+  .add.f32` in the order they land, as FlashAttention sums them. It skips the
+  rounding and is not deterministic. P then carries no grid exponent, whose
+  subtraction rounds the softmax's argument, so dK and dV move in their last
+  bits too.
+- `persistent=False`: a dense call runs one CTA per work tile under
+  `SingleTileScheduler` instead of the persistent work list, with no
+  canonical dK/dV records.
+"""
+
+from __future__ import annotations
+
+from .launch import BackwardLaunch, _prepare
+
+
+def prepare_backward_variant(
+    q, k, v, o, do, lse, *, causal, fp32_dq=False, persistent=True, **kwargs
+) -> BackwardLaunch:
+    """`prepare_backward` with dQ summed as fp32 (`fp32_dq`) and the dense
+    scheduler `persistent`; the other keyword arguments are
+    `prepare_backward`'s. The defaults are the shipped kernel."""
+    return _prepare(
+        q, k, v, o, do, lse, causal=causal, dq_fp32=fp32_dq, persistent=persistent, **kwargs
+    )

fold_attention-0.1.0/src/fold_attention/backward/grid.py

@@ -0,0 +1,100 @@
+"""The backward's integer grids, derived from bounds the kernels compute.
+
+Every kernel that quantises or reconstructs a partial evaluates the same bound
+with the same rounded instructions, so each arrives at the same scale. The
+inputs are the preprocess's maxima (`red.max` over float bits), which do not
+depend on the order CTAs finish in or on anything outside the request.
+
+dQ, per (request, KV head): `|sum_j dS_ij K_jd| <= (sum_j |dS_ij|) max|K|`
+and, with P summing to one, `sum_j |dS_ij| <= |dO_i| max_j |V_j| + |D_i|`,
+where `|V_j| <= C_D max|V|` and `C_D` is sqrt(D) rounded up to a power of two.
+dV, on the canonical-record path: `|dV_jd| <= C max_i |dO_id|`, the largest
+element of dO, since the bound is per output component. dK there:
+`|dK_j| <= C (max|dO| C_D max|V| + max|D|) max|Q|`. `C` bounds the column
+mass `sum_{h,i} P_hij`; the only data-free bound is the number of query rows
+`G n`, which is why dK needs a 61-bit grid. Every bound is widened by 2^-6 for
+the bf16 rounding of P and dS.
+"""
+
+import math
+
+import cutlass
+from cutlass import Float32, Int32
+
+from ..ptx import f32_rn
+
+# stats[b, h_kv, i], float bits, zero before the call: max |K|, max |V|, max
+# ||dO_i|| over the group's rows, max |D_i| over them, and for the record
+# path max |dO_id| over every element and max |Q| over the group
+ST_K, ST_V, ST_DO, ST_DELTA, ST_DOC, ST_Q = 0, 1, 2, 3, 4, 5
+N_STATS = 8
+
+DQ_BITS = 30
+DKV_BITS = 61
+
+
+def root_d(D):
+    """sqrt(D) rounded up to a power of two: exact, and a bound."""
+    return float(2 ** math.ceil(math.log2(math.sqrt(D))))
+
+
+def pow2_scale(bound, bits: int) -> Float32:
+    """`2^(bits - ceil(log2 (bound (1 + 2^-6))))`, clamped to [2^-100, 2^100].
+    A zero bound takes 2^100, since its partials are all zero."""
+    b = f32_rn("mul.rn.f32", bound, Float32(1.0 + 2.0**-6))
+    bb = b.bitcast(Int32)
+    clog = (
+        ((bb >> 23) & 0xFF) - 127 + Int32(cutlass.select_((bb & 0x7FFFFF) != 0, Int32(1), Int32(0)))
+    )
+    sexp = Int32(bits) - clog
+    sexp = Int32(cutlass.select_(b > Float32(0.0), sexp, Int32(100)))
+    sexp = Int32(cutlass.select_(sexp > 100, Int32(100), sexp))
+    sexp = Int32(cutlass.select_(sexp < -100, Int32(-100), sexp))
+    return ((sexp + 127) << 23).bitcast(Float32)
+
+
+def inv_pow2(scale) -> Float32:
+    """`1 / scale` for a power of two, from its bits."""
+    return (Int32(0x7F000000) - scale.bitcast(Int32)).bitcast(Float32)
+
+
+def log2_pow2(scale) -> Float32:
+    """The exponent of a power of two, as a float."""
+    return ((scale.bitcast(Int32) >> 23) - Int32(127)).to(Float32)
+
+
+def load_stats(mStats, batch_idx, head_kv, which=range(N_STATS)):
+    """The maxima `which` of one (request, KV head), by index."""
+    return {i: mStats[batch_idx, head_kv, i].bitcast(Float32) for i in which}
+
+
+def dq_scale(mStats, batch_idx, head_kv, cd: float) -> Float32:
+    """dQ's scale for one (request, KV head).
+
+    A constant exponent is free here: `P 2^s` is `exp2(c S - lse + s)`, so
+    the softmax's own exponent carries it and the dQ partials leave the GEMM
+    on the grid already."""
+    st = load_stats(mStats, batch_idx, head_kv, (ST_K, ST_V, ST_DO, ST_DELTA))
+    s = f32_rn(
+        "add.rn.f32",
+        f32_rn("mul.rn.f32", st[ST_DO], f32_rn("mul.rn.f32", st[ST_V], Float32(cd))),
+        st[ST_DELTA],
+    )
+    return pow2_scale(f32_rn("mul.rn.f32", s, st[ST_K]), DQ_BITS)
+
+
+def dv_scale(st, rows, bits: int) -> Float32:
+    """dV's scale from `load_stats`, for `rows` query rows bounding the column
+    mass."""
+    return pow2_scale(f32_rn("mul.rn.f32", rows, st[ST_DOC]), bits)
+
+
+def dk_scale(st, rows, cd: float, bits: int) -> Float32:
+    """dK's scale from `load_stats`, for `rows` query rows bounding the column
+    mass."""
+    s = f32_rn(
+        "add.rn.f32",
+        f32_rn("mul.rn.f32", st[ST_DO], f32_rn("mul.rn.f32", st[ST_V], Float32(cd))),
+        st[ST_DELTA],
+    )
+    return pow2_scale(f32_rn("mul.rn.f32", rows, f32_rn("mul.rn.f32", s, st[ST_Q])), bits)