stream-cqsa 2.1.2__py3-none-any.whl

stream_cqsa/__init__.py

@@ -0,0 +1,90 @@
+__version__ = "2.1.2"
+
+from .stable_stream import (
+    stream_cqsa_forward,
+    stream_cqsa_backward,
+    TraceRecorder,
+)
+from .oom_fallback import (stream_cqsa_auto, attention_oom_safe,
+                           ESCALATION, ESCALATION_FAST)
+from .native_autograd import stream_cqsa_attn, StreamCQSAAttention
+from .api import attention, estimate, patch_sdpa, unpatch_sdpa, patched_sdpa, kernels_available
+from .doctor import doctor
+from .progress import verbose_enabled
+
+from .interface import (
+    flash_attn_func,
+    flash_attn_func_cqs,
+    flash_attn_func_cqs_group_bits,
+    flash_attn_bwd_cqs_group_bits,
+    flash_attn_func_cqsa,
+)
+from .config import (
+    DEFAULT_PROBING_HISTORY,
+    DEFAULT_PROBING_HISTORY_DIR,
+    GLOBAL_STREAM_SEED,
+    MEM_BUDGET_COE,
+)
+from .cqsa_probe_capacity import CQSAProbe
+from .cqs_attention_stream import (
+    CQSAStreamRunner,
+    SubsequenceAttentionFn,
+    compare_O,
+    run_scqsa_stream,
+    stream_cqsa,
+)
+from .attention_kernel import (
+    default_subsequence_attention,
+    custom_attn,
+)
+from .cqs_mask import CQS_mask
+from .memory_fitting import fitting_mem, pred_mem, best_seq_length, estimate_memory_time_from_model
+from .autograd_op import stream_cqsa_autograd
+
+__all__ = [
+    # --- the one-call API -------------------------------------------------
+    "attention",               # SDPA signature; monolithic when it fits, exact decomposition when it does not
+    "estimate",                # dry run: what a call would cost here
+    "patch_sdpa", "unpatch_sdpa", "patched_sdpa",
+    "doctor", "kernels_available",
+    # --- primary entry points -------------------------------------------
+    "stream_cqsa_auto",        # "just run it": escalates until it fits
+    "stream_cqsa_forward",     # explicit forward, returns (out, lse)
+    "stream_cqsa_backward",    # explicit backward, needs the global lse
+    "stream_cqsa_attn",        # autograd-aware; supports .backward()
+    "StreamCQSAAttention",     # the same, as an nn.Module
+    "attention_oom_safe",
+    "ESCALATION",
+    "ESCALATION_FAST",
+    "TraceRecorder",
+    # --- everything else --------------------------------------------------
+    "flash_attn_func",
+    "flash_attn_func_cqs",
+    "flash_attn_func_cqs_group_bits",
+    "flash_attn_bwd_cqs_group_bits",
+    "flash_attn_func_cqsa",
+    "CQSAProbe",
+    "CQSAStreamRunner",
+    "SubsequenceAttentionFn",
+    "DEFAULT_PROBING_HISTORY",
+    "DEFAULT_PROBING_HISTORY_DIR",
+    "GLOBAL_STREAM_SEED",
+    "MEM_BUDGET_COE",
+    "run_scqsa_stream",
+    "stream_cqsa",
+    "compare_O",
+    "default_subsequence_attention",
+    "custom_attn",
+    "CQS_mask",
+    "fitting_mem",
+    "pred_mem",
+    "best_seq_length",
+    "estimate_memory_time_from_model",
+    "stream_cqsa_autograd",
+]
+
+# v2 additions
+from .autoconfig import (detect_hardware, hardware_from_dict, plan, calibrate, autotune,  # noqa: E402
+                         auto_attention, QUORUM_SETS, CostModel, Plan)
+from .devkit import compare_kernels, quick_bench, Config, run_config  # noqa: E402
+from .adapters import flex_inner, dense_inner  # noqa: E402

stream_cqsa/adapters.py

@@ -0,0 +1,127 @@
+"""
+Adapters that turn a monolithic attention implementation into a Stream-CQSA
+inner kernel without writing CUDA.
+
+What an inner kernel must provide (see devkit.compare_kernels):
+
+    inner(q_i, k_i, v_i, group_bits, *, causal, scale, **_) -> (out_i, lse_i)
+
+i.e. attention restricted to the pair set {(r, c): bits[r] & bits[c] == 0,
+(c <= r if causal)}, returning the normalised output ``out_i [B, L, H, D]``
+AND the per-row log-sum-exp of the kept scores ``lse_i [B, L, H]`` (both
+token-major, float32). Any implementation that (a) accepts an
+arbitrary pairwise mask and (b) returns the log-sum-exp can therefore be
+plugged in mechanically. Two such families exist in PyTorch today:
+
+* ``flex_inner(score_mod=None)`` -- torch's FlexAttention. The CQS pair set is
+  expressed as a ``mask_mod`` (a block mask is built per subproblem, so the
+  masked tiles are skipped by the compiled kernel too), and ``return_lse=True``
+  gives the lse. Any attention that FlexAttention can express -- ALiBi,
+  soft-capping, sliding windows, document masks, relative biases -- becomes a
+  Stream-CQSA inner kernel by passing its ``score_mod``/extra mask, with no
+  other change. That is the "automatic conversion" for the score-modification
+  family.
+
+* ``dense_inner(mono_masked)`` -- for small subproblems, any callable that
+  takes a dense boolean mask and returns (out, lse) (or just out, in which
+  case the lse is recomputed densely). This is the differential-testing route.
+
+Kernels that return NO log-sum-exp cannot be adapted losslessly: the merge of
+subproblems needs the row statistics, and recovering them from the output
+alone is impossible (the output is normalised). For those the only exact
+option is a second pass that computes the lse (``lse_from_scores``), which
+costs another QK^T; the devkit reports that overhead honestly.
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import torch
+
+_flex_cache: dict = {}
+
+
+def _bits_pair_mask(bits: torch.Tensor, causal: bool):
+    def mask_mod(b, h, q_idx, kv_idx):
+        keep = (bits[q_idx] & bits[kv_idx]) == 0
+        if causal:
+            keep = keep & (kv_idx <= q_idx)
+        return keep
+    return mask_mod
+
+
+def flex_inner(score_mod: Callable | None = None, *, extra_mask_mod: Callable | None = None,
+               block_size: int = 128, compile: bool = True, global_positions: bool = True) -> Callable:
+    """
+    Build a Stream-CQSA inner kernel from torch FlexAttention.
+
+    `score_mod(score, b, h, q_idx, kv_idx)` and `extra_mask_mod(b, h, q_idx, kv_idx)`
+    are written against GLOBAL token positions (as in the monolithic call): a
+    subproblem sees a gathered subsequence, so the engine hands the inner kernel
+    the gather index (`token_ids`) and the adapter remaps local indices through
+    it before calling the user's functions. That is the one thing an automatic
+    conversion must get right for position-dependent attention (ALiBi, sliding
+    windows, document masks); with `global_positions=False` the raw local
+    indices are passed instead (only correct for position-free score_mods).
+    """
+    from torch.nn.attention.flex_attention import flex_attention, create_block_mask
+    fa = torch.compile(flex_attention, dynamic=False) if compile else flex_attention
+
+    def inner(q_i, k_i, v_i, group_bits, *, causal: bool, scale: float, token_ids=None, **_):
+        B, L, H, D = q_i.shape
+        dev = q_i.device
+        bits = group_bits.to(dev, torch.int64)
+        mm = _bits_pair_mask(bits, causal)
+        ids = token_ids.to(dev, torch.int64) if (token_ids is not None and global_positions) else None
+        if extra_mask_mod is not None:
+            base = mm
+            if ids is not None:
+                mm = lambda b, h, qi, ki: base(b, h, qi, ki) & extra_mask_mod(b, h, ids[qi], ids[ki])
+            else:
+                mm = lambda b, h, qi, ki: base(b, h, qi, ki) & extra_mask_mod(b, h, qi, ki)
+        sm = score_mod
+        if score_mod is not None and ids is not None:
+            sm = lambda score, b, h, qi, ki: score_mod(score, b, h, ids[qi], ids[ki])
+        bm = create_block_mask(mm, B=None, H=None, Q_LEN=L, KV_LEN=L, device=dev, BLOCK_SIZE=block_size)
+        qb, kb, vb = (t.transpose(1, 2) for t in (q_i, k_i, v_i))     # [B,H,L,D]
+        out, lse = fa(qb, kb, vb, score_mod=sm, block_mask=bm, scale=float(scale), return_lse=True)
+        # rows with no kept pair: flex returns lse=-inf and out=0, which is the contract.
+        # Engine contract: out_i [B,L,H,D], lse_i [B,L,H] (token-major, like the CUDA kernel's return).
+        return out.transpose(1, 2).float(), lse.float().transpose(1, 2)
+
+    return inner
+
+
+def lse_from_scores(q_i, k_i, v_i, keep: torch.Tensor, scale: float) -> torch.Tensor:
+    """Dense fp32 per-row lse over kept pairs (second pass for kernels without lse). [B,H,L]"""
+    s = torch.matmul(q_i.transpose(1, 2).float(), k_i.transpose(1, 2).float().transpose(-1, -2)) * float(scale)
+    s = s.masked_fill(~keep[None, None], float("-inf"))
+    return torch.logsumexp(s, dim=-1)
+
+
+def dense_inner(mono_masked: Callable, *, returns_lse: bool = True) -> Callable:
+    """
+    Wrap ``mono_masked(q, k, v, mask, scale) -> (out [B,H,L,D], lse [B,H,L])`` (or
+    just out if returns_lse=False) where mask is a dense bool [L, L] (True = keep).
+    Only for small L (the mask is O(L^2)).
+    """
+    def inner(q_i, k_i, v_i, group_bits, *, causal: bool, scale: float, **_):
+        L = q_i.shape[1]
+        bits = group_bits.to(q_i.device, torch.int64)
+        keep = (bits[:, None] & bits[None, :]) == 0
+        if causal:
+            keep = keep & torch.ones(L, L, dtype=torch.bool, device=q_i.device).tril()
+        qb, kb, vb = (t.transpose(1, 2) for t in (q_i, k_i, v_i))
+        r = mono_masked(qb, kb, vb, keep, float(scale))
+        if returns_lse:
+            out, lse = r
+        else:
+            out, lse = r, lse_from_scores(q_i, k_i, v_i, keep, scale)
+        return out.transpose(1, 2).float(), lse.float().transpose(1, 2)   # [B,L,H,D], [B,L,H]
+    return inner
+
+
+def sdpa_masked(qb, kb, vb, keep, scale):
+    """Reference monolithic-with-mask kernel (torch SDPA, no lse): use with dense_inner(..., returns_lse=False)."""
+    import torch.nn.functional as F
+    return F.scaled_dot_product_attention(qb, kb, vb, attn_mask=keep, scale=scale)

stream_cqsa/api.py

@@ -0,0 +1,363 @@
+"""
+The one entry point: ``stream_cqsa.attention`` with the signature of
+``torch.nn.functional.scaled_dot_product_attention``.
+
+    out = stream_cqsa.attention(q, k, v, is_causal=True)
+
+* below the memory boundary it IS the monolithic call (SDPA / FlashAttention);
+* above it, the planner picks the decomposition (depth, quorum set, accumulator
+  placement, host streaming, concurrency) from the free memory and the call
+  runs exactly, on the best kernel available (native wave kernel, CUDA
+  extension, or the Triton kernels when nothing is compiled);
+* gradients flow when the inputs require them;
+* ``verbose=True`` (or ``CQSA_VERBOSE=1``) shows what is happening.
+
+``patch_sdpa()`` swaps this in for ``F.scaled_dot_product_attention`` so an
+existing model runs unchanged; ``estimate(N, ...)`` is the dry run that says
+what a call would cost before it is made.
+"""
+from __future__ import annotations
+
+import contextlib
+import os
+import warnings
+from typing import Any, Optional
+
+import torch
+import torch.nn.functional as F
+
+from .progress import verbose_enabled
+
+__all__ = ["attention", "estimate", "patch_sdpa", "unpatch_sdpa", "patched_sdpa", "kernels_available"]
+
+
+def kernels_available() -> dict[str, bool]:
+    """Which attention kernels this environment can run."""
+    from . import interface as I
+    out = {"cuda_extension": I.cqsa_cuda is not None, "cuda_extension_noncausal": I.cqsa_cuda_noncausal is not None}
+    try:
+        from .native_wave import native_available
+        out["native_wave"] = native_available()
+    except Exception:
+        out["native_wave"] = False
+    try:
+        import triton  # noqa: F401
+        from . import triton_kernel  # noqa: F401
+        out["triton"] = True
+    except Exception:
+        out["triton"] = False
+    try:
+        import flash_attn  # noqa: F401
+        out["flash_attn"] = True
+    except Exception:
+        out["flash_attn"] = False
+    return out
+
+
+def _check_inputs(q, k, v, attn_mask, dropout_p, enable_gqa):
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(f"stream_cqsa.attention expects q/k/v as [B, H, N, D]; got {tuple(q.shape)}, {tuple(k.shape)}, {tuple(v.shape)}")
+    if attn_mask is not None:
+        raise NotImplementedError("stream_cqsa.attention supports the plain and the causal (is_causal=True) cases; an explicit "
+                                  "attn_mask is not supported. For FlexAttention-style masks use stream_cqsa.adapters.flex_inner.")
+    if dropout_p:
+        raise NotImplementedError("stream_cqsa.attention does not support dropout (dropout_p must be 0).")
+    if enable_gqa or k.shape[1] != q.shape[1]:
+        raise NotImplementedError("stream_cqsa.attention needs the same number of heads for q and k/v (no GQA); repeat k/v heads first.")
+    if q.shape != k.shape or q.shape != v.shape:
+        raise ValueError("q, k and v must have the same shape [B, H, N, D] (cross-attention with N_q != N_k is not supported)")
+    if q.dtype not in (torch.float16, torch.bfloat16):
+        raise TypeError(f"stream_cqsa.attention needs fp16 or bf16 inputs (got {q.dtype}); cast with .half() or .bfloat16(). "
+                        "Accumulation is fp32 internally regardless.")
+    if not (q.dtype == k.dtype == v.dtype):
+        raise TypeError("q, k and v must share a dtype")
+    if not (q.device == k.device == v.device):
+        raise ValueError("q, k and v must be on the same device")
+
+
+def _is_oom(exc: BaseException) -> bool:
+    return isinstance(exc, torch.cuda.OutOfMemoryError) or "out of memory" in str(exc).lower()
+
+
+def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask=None, dropout_p: float = 0.0,
+              is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, *,
+              verbose: Optional[bool] = None, hardware=None, kernel: str = "auto", plan_only: bool = False,
+              return_plan: bool = False, **overrides) -> torch.Tensor:
+    """
+    Exact attention that always fits. Drop-in for ``F.scaled_dot_product_attention``
+    (same positional signature; masks, dropout and GQA are rejected with a message).
+
+    q/k/v ``[B, H, N, D]`` fp16/bf16, on a device or in host memory. Returns the
+    output in the inputs' dtype and device. Differentiable when the inputs require
+    gradients. ``kernel``: "auto" (native wave kernel when built and applicable,
+    else the CUDA extension, else Triton), "wave", "cuda", "triton".
+    ``overrides`` (itr=, c=, interest_set=, max_parallel=, ...) pin engine settings.
+    ``plan_only=True`` returns the plan without computing anything.
+    """
+    _INSIDE[0] += 1
+    try:
+        return _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, verbose=verbose, hardware=hardware,
+                          kernel=kernel, plan_only=plan_only, return_plan=return_plan, **overrides)
+    finally:
+        _INSIDE[0] -= 1
+
+
+def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, verbose, hardware, kernel, plan_only,
+               return_plan, **overrides):
+    from .autoconfig import plan as _plan, hardware_from_dict, detect_hardware
+    _check_inputs(q, k, v, attn_mask, dropout_p, enable_gqa)
+    B, H, N, D = q.shape
+    if scale is None:
+        scale = float(D) ** -0.5
+    on_cuda = q.device.type == "cuda"
+    out_device = q.device
+    hw = hardware_from_dict(hardware) if isinstance(hardware, dict) else (hardware or detect_hardware())
+    direction = "bwd" if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad) else "fwd"
+    p = _plan(N=N, B=B, H=H, D=D, dtype=q.dtype, causal=bool(is_causal), hardware=hw, direction=direction,
+              allow_distributed=False)
+    if plan_only:
+        return p
+    vb = verbose_enabled(verbose)
+
+    # ---- monolithic: the call fits, so it IS the normal kernel ------------------
+    if p.mode == "mono" and not overrides:
+        if vb:
+            print(f"Stream-CQSA: N={N} fits ({p.est_peak_gib:.1f} GiB) -- monolithic call, no decomposition", flush=True)
+        qd, kd, vd = ((t if on_cuda else t.to("cuda", non_blocking=True)) for t in (q, k, v))
+        try:
+            out = _TRUE_SDPA(qd, kd, vd, is_causal=bool(is_causal), scale=scale)
+            out = out if on_cuda else out.to(q.device)
+            return (out, p) if return_plan else out
+        except Exception as exc:                                     # noqa: BLE001
+            if not _is_oom(exc):
+                raise
+            if vb:
+                print("Stream-CQSA: the monolithic call ran out of memory after all; decomposing", flush=True)
+            del qd, kd, vd
+            torch.cuda.empty_cache()
+            p = _plan(N=N, B=B, H=H, D=D, dtype=q.dtype, causal=bool(is_causal), hardware=detect_hardware(),
+                      direction=direction, allow_distributed=False)
+            if p.mode == "mono":     # the planner still thinks it fits: force a decomposition
+                p.mode, p.itr, p.c, p.interest_set = "cqsa", 1, 7, (0, 1, 3)
+
+    # ---- decomposed ------------------------------------------------------------
+    kw = p.engine_kwargs() if p.mode != "mono" else dict(itr=1, c=7, interest_set=(0, 1, 3))
+    if on_cuda and kw.get("stream_from_host"):
+        # The inputs are on the device. Prefer the fastest configuration that leaves them
+        # there; stream from the host only if nothing device-resident fits (then a pinned
+        # host copy is made -- the caller's device tensors stay where they are).
+        dev_ok = [c for c in p.candidates if c["ok"] and c["mode"] == "cqsa" and not c["stream_from_host"]]
+        if dev_ok:
+            c0 = min(dev_ok, key=lambda c: c["time"])
+            kw = dict(itr=int(c0["itr"]), c=int(c0["c"]), interest_set=tuple(c0["interest_set"]),
+                      low_memory=(c0["acc"] == "cpu"), accumulate_on_gpu=(c0["acc"] == "gpu"),
+                      stream_from_host=False, max_parallel=int(c0["n_par"]), shared_chunks=False)
+        elif direction == "bwd":
+            # Copying the inputs to the host would sever the autograd graph. Keep them on the
+            # device, park the accumulator in host memory, and let the engine escalate the
+            # depth on its own if the device is still short.
+            if vb:
+                print("Stream-CQSA: no device-resident configuration fits the estimate; keeping Q/K/V on the device with a "
+                      "host accumulator and escalating on demand", flush=True)
+            kw.update(stream_from_host=False, low_memory=True, accumulate_on_gpu=False, shared_chunks=False)
+        else:
+            if vb:
+                print("Stream-CQSA: no device-resident configuration fits; copying Q/K/V to pinned host memory and streaming", flush=True)
+            if direction == "bwd":
+                q, k, v = (t.to("cpu") for t in (q, k, v))                 # differentiable copies: the graph stays intact
+            else:
+                q, k, v = (t.detach().to("cpu").pin_memory() for t in (q, k, v))
+            on_cuda = False
+    kw.update(overrides)
+    if not kw.get("stream_from_host"):
+        kw["shared_chunks"] = False
+    use_wave = _pick_wave(kernel, kw, is_causal, q, direction)
+    if vb:
+        print(f"Stream-CQSA: {p.reason}", flush=True)
+    if use_wave:
+        from .native_wave import wave_attention, wave_forward
+        wkw = dict(itr=int(kw.get("itr", 1)), c=int(kw.get("c", 7)), interest_set=tuple(kw.get("interest_set", (0, 1, 3))))
+        try:
+            if direction == "bwd":
+                out = wave_attention(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, **wkw)
+            else:
+                out, _ = wave_forward(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, **wkw)
+        except Exception as exc:                                     # noqa: BLE001
+            if not _is_oom(exc) or kernel == "wave":
+                raise
+            # the wave engine keeps its accumulator on the device; the classic engine can offload it
+            if vb:
+                print("Stream-CQSA: the wave engine ran out of memory; retrying on the streaming engine", flush=True)
+            torch.cuda.empty_cache()
+            use_wave = False
+    if not use_wave:
+        from .native_autograd import stream_cqsa_attn
+        from .stable_stream import stream_cqsa_forward
+        if kernel == "triton":
+            os.environ["CQSA_FORWARD"] = "triton"; os.environ["CQSA_BACKWARD"] = "triton"
+        host = not on_cuda
+        if direction == "bwd":
+            out = stream_cqsa_attn(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose,
+                                   itr=kw.get("itr", "auto"), c=kw.get("c", 7), interest_set=kw.get("interest_set", (0, 1, 3)),
+                                   stream_from_host=bool(kw.get("stream_from_host", host)),
+                                   accumulate_on_gpu=bool(kw.get("accumulate_on_gpu", True)),
+                                   max_parallel=kw.get("max_parallel"))
+        else:
+            kw.setdefault("stream_from_host", host)
+            # leave the fp32 result where the accumulator is; _deliver casts and moves it
+            out, _ = stream_cqsa_forward(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, out_device="acc", **kw)
+    out = _deliver(out, out_device, q.dtype, vb)
+    return (out, p) if return_plan else out
+
+
+def _deliver(out: torch.Tensor, device, dtype, vb: bool) -> torch.Tensor:
+    """Hand the result over in the caller's dtype and on the caller's device.
+
+    The engines return fp32; the cast needs room for a second copy. When the device
+    cannot take it, the cast is done in host memory and the result moved back; if
+    even the half-size result does not fit next to what the caller keeps on the
+    device, it is returned in host memory with a warning rather than failing after
+    the work is done."""
+    if out.dtype == dtype and out.device == device:
+        return out
+    try:
+        return out.to(device=device, dtype=dtype)
+    except Exception as exc:                                         # noqa: BLE001
+        if not _is_oom(exc):
+            raise
+    host = out.to("cpu") if out.device.type != "cpu" else out
+    del out
+    torch.cuda.empty_cache()
+    host = host.to(dtype)
+    if device.type == "cpu":
+        return host
+    try:
+        return host.to(device)
+    except Exception as exc:                                         # noqa: BLE001
+        if not _is_oom(exc):
+            raise
+    warnings.warn(f"Stream-CQSA: the result ({host.numel() * host.element_size() / 2**30:.1f} GiB) does not fit on {device} "
+                  f"next to what is already there; returning it in host memory.", RuntimeWarning, stacklevel=3)
+    return host
+
+
+def _pick_wave(kernel: str, kw: dict, is_causal: bool, q: torch.Tensor, direction: str) -> bool:
+    if kernel == "wave":
+        return True
+    if kernel in ("cuda", "triton"):
+        return False
+    if kernel != "auto":
+        raise ValueError(f"kernel must be 'auto', 'wave', 'cuda' or 'triton' (got {kernel!r})")
+    try:
+        from .native_wave import native_available
+        if not native_available():
+            return False
+    except Exception:
+        return False
+    # the wave engine keeps its accumulator (and, in the backward, the fp32 gradients) on the device
+    if kw.get("low_memory") or kw.get("accumulate_on_gpu") is False:
+        return False
+    from .native_wave import native_supports
+    return native_supports(q.dtype, int(q.shape[-1]))      # depends on the kernel set the extension was built with
+
+
+# ---------------------------------------------------------------------------
+# dry run
+# ---------------------------------------------------------------------------
+def estimate(N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, causal: bool = True,
+             hardware=None, direction: str = "both", print_table: bool = True) -> dict[str, Any]:
+    """
+    What would ``attention`` do for this shape on this machine, and what would it
+    cost? Prints (and returns) the monolithic estimate, the chosen decomposition
+    and the runner-up configurations, for the forward and the backward.
+    """
+    from .autoconfig import plan as _plan, hardware_from_dict, detect_hardware, _cname, GIB
+    hw = hardware_from_dict(hardware) if isinstance(hardware, dict) else (hardware or detect_hardware())
+    res: dict[str, Any] = {"N": N, "B": B, "H": H, "D": D, "dtype": str(dtype), "causal": causal, "hardware": hw.summary()}
+    dirs = ("fwd", "bwd") if direction == "both" else (direction,)
+    lines = [f"Stream-CQSA estimate for N={N:,} B={B} H={H} D={D} {str(dtype).replace('torch.', '')} {'causal' if causal else 'non-causal'}",
+             f"  hardware: {hw.summary()}"]
+    itemsize = torch.tensor([], dtype=dtype).element_size()
+    for d in dirs:
+        p = _plan(N=N, B=B, H=H, D=D, dtype=dtype, causal=causal, hardware=hw, direction=d, allow_distributed=False)
+        mono = p.candidates[0]
+        feas = sorted([c for c in p.candidates if c["ok"] and c["mode"] != "mono"], key=lambda c: c["time"])[:3]
+        r = dict(plan=p.name(), est_time_s=p.est_time_s, est_peak_gib=p.est_peak_gib, est_host_gib=p.est_host_gib,
+                 mono_fits=bool(mono["ok"]), mono_peak_gib=mono["peak"], mono_time_s=mono["time"], reason=p.reason,
+                 runners_up=[dict(config=_cname(c), time_s=c["time"], peak_gib=c["peak"], host_gib=c["host"]) for c in feas],
+                 feasible=any(c["ok"] for c in p.candidates))
+        res[d] = r
+        label = "forward" if d == "fwd" else "backward (fwd+bwd step)"
+        lines.append(f"  {label}:")
+        lines.append(f"    monolithic FlashAttention: {'fits' if mono['ok'] else 'does NOT fit'} ({mono['peak']:.1f} GiB), ~{mono['time']:.2f} s")
+        if p.mode == "mono":
+            lines.append(f"    -> monolithic call")
+        elif r["feasible"]:
+            lines.append(f"    -> Stream-CQSA {p.name()}: ~{p.est_time_s:.1f} s, {p.est_peak_gib:.1f} GiB on the device"
+                         + (f", {p.est_host_gib:.0f} GiB host" if p.est_host_gib else ""))
+            for c in feas[1:]:
+                lines.append(f"       also: {_cname(c)} ~{c['time']:.1f} s, {c['peak']:.1f} GiB")
+        else:
+            floor = (4 + 4) * N * B * H * D / GIB
+            lines.append(f"    -> nothing fits: the fp32 output + accumulator alone are {floor:.1f} GiB on the device; "
+                         f"reduce B/H/D, add host memory, or use more devices")
+    if print_table:
+        print("\n".join(lines), flush=True)
+    res["text"] = "\n".join(lines)
+    return res
+
+
+# ---------------------------------------------------------------------------
+# monkeypatch F.scaled_dot_product_attention
+# ---------------------------------------------------------------------------
+_ORIGINAL_SDPA = None
+_PATCH_MIN_TOKENS = 65536
+_TRUE_SDPA = F.scaled_dot_product_attention          # captured at import, before any patch
+_INSIDE = [0]                                          # re-entrancy: calls made by the engines go to the original
+
+
+def _patched_sdpa(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False):
+    orig = _ORIGINAL_SDPA or _TRUE_SDPA
+    supported = (attn_mask is None and not dropout_p and not enable_gqa and query.dim() == 4
+                 and query.dtype in (torch.float16, torch.bfloat16) and key.shape == query.shape == value.shape
+                 and query.shape[-2] >= _PATCH_MIN_TOKENS and _INSIDE[0] == 0)
+    if not supported:
+        return orig(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale,
+                    **({"enable_gqa": enable_gqa} if enable_gqa else {}))
+    return attention(query, key, value, is_causal=is_causal, scale=scale)
+
+
+def patch_sdpa(min_tokens: int = 65536) -> None:
+    """
+    Route ``torch.nn.functional.scaled_dot_product_attention`` through
+    ``stream_cqsa.attention`` for the calls it handles (4-D fp16/bf16, no mask,
+    no dropout, no GQA, N >= ``min_tokens``); everything else goes to the original.
+    Below the memory boundary the routed call is still the monolithic kernel, so
+    the patch changes nothing until a call would have run out of memory.
+    """
+    global _ORIGINAL_SDPA, _PATCH_MIN_TOKENS
+    _PATCH_MIN_TOKENS = int(min_tokens)
+    if _ORIGINAL_SDPA is None:
+        _ORIGINAL_SDPA = F.scaled_dot_product_attention
+        F.scaled_dot_product_attention = _patched_sdpa
+        torch.nn.functional.scaled_dot_product_attention = _patched_sdpa
+
+
+def unpatch_sdpa() -> None:
+    global _ORIGINAL_SDPA
+    if _ORIGINAL_SDPA is not None:
+        F.scaled_dot_product_attention = _ORIGINAL_SDPA
+        torch.nn.functional.scaled_dot_product_attention = _ORIGINAL_SDPA
+        _ORIGINAL_SDPA = None
+
+
+@contextlib.contextmanager
+def patched_sdpa(min_tokens: int = 65536):
+    """``with stream_cqsa.patched_sdpa(): model(...)``"""
+    was = _ORIGINAL_SDPA is not None
+    patch_sdpa(min_tokens)
+    try:
+        yield
+    finally:
+        if not was:
+            unpatch_sdpa()

stream_cqsa/attention_kernel/Custom.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Sequence, Tuple
+
+import torch
+
+
+def _group_bits_to_dense_mask(
+    *,
+    group_bits: torch.Tensor,
+    device: torch.device,
+) -> torch.Tensor:
+    """
+    Build dense local mask [L, L] from per-token int64 group bits.
+    True means masked.
+    """
+    bits = group_bits.to(device=device, dtype=torch.int64).contiguous()
+    bits_row = bits.view(bits.numel(), 1)
+    bits_col = bits.view(1, bits.numel())
+    return torch.bitwise_and(bits_row, bits_col).ne(0)
+
+
+def _group_runs_to_dense_mask(
+    *,
+    local_size: int,
+    group_runs: Sequence[Sequence[Tuple[int, int]]],
+    device: torch.device,
+) -> torch.Tensor:
+    """
+    Build dense local mask [L, L] where True means masked.
+    """
+    mask = torch.zeros((int(local_size), int(local_size)), dtype=torch.bool, device=device)
+    for runs in group_runs:
+        idx_parts: List[torch.Tensor] = []
+        for s, e in runs:
+            si, ei = int(s), int(e)
+            if ei > si:
+                idx_parts.append(torch.arange(si, ei, device=device, dtype=torch.long))
+        if len(idx_parts) == 0:
+            continue
+        idx = idx_parts[0] if len(idx_parts) == 1 else torch.cat(idx_parts, dim=0)
+        mask[idx.unsqueeze(1), idx.unsqueeze(0)] = True
+    return mask
+
+
+def custom_attn(
+    *,
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cqs_mask: Dict[str, Any],
+    softmax_scale: float,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Naive reference subsequence attention plugin.
+
+    Steps:
+    1) R_i = (Q @ K^T) * scale
+    2) apply CQS mask M_i on R_i (masked positions -> -inf)
+    3) P_i = exp(R_i)   (no softmax normalization)
+    4) Den_i = row_sum(P_i)
+    5) Num_i = P_i @ V
+
+    Returns:
+    - Num_i: [B, L, H, D] float32
+    - Den_i: [B, H, L] float32
+    """
+    local_size = int(cqs_mask["local_size"])
+    group_bits_cpu = cqs_mask.get("group_bits_cpu", None)
+    if isinstance(group_bits_cpu, torch.Tensor):
+        M_i = _group_bits_to_dense_mask(
+            group_bits=group_bits_cpu,
+            device=q.device,
+        )
+    else:
+        M_i = _group_runs_to_dense_mask(
+            local_size=local_size,
+            group_runs=cqs_mask["group_runs"],
+            device=q.device,
+        )  # [L, L], True means masked
+
+    q_bhld = q.transpose(1, 2).float()  # [B, H, L, D]
+    k_bhld = k.transpose(1, 2).float()  # [B, H, L, D]
+    v_bhld = v.transpose(1, 2).float()  # [B, H, L, D]
+
+    R_i = torch.matmul(q_bhld, k_bhld.transpose(-2, -1)) * float(softmax_scale)  # [B, H, L, L]
+    if bool(M_i.any().item()):
+        R_i = R_i.masked_fill(M_i.unsqueeze(0).unsqueeze(0), float("-inf"))
+
+    P_i = torch.exp(R_i)  # [B, H, L, L]
+    P_i = torch.nan_to_num(P_i, nan=0.0, posinf=0.0, neginf=0.0)
+
+    Den_i = P_i.sum(dim=-1)  # [B, H, L]
+    Num_i = torch.matmul(P_i, v_bhld).transpose(1, 2).contiguous()  # [B, L, H, D]
+    Num_i = torch.nan_to_num(Num_i, nan=0.0, posinf=0.0, neginf=0.0)
+    Den_i = torch.nan_to_num(Den_i, nan=0.0, posinf=0.0, neginf=0.0)
+    return Num_i, Den_i