stream-cqsa 2.2.2__tar.gz → 2.2.3__tar.gz

{stream_cqsa-2.2.2/stream_cqsa.egg-info → stream_cqsa-2.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-cqsa
-Version: 2.2.2
+Version: 2.2.3
 Summary: Stream-CQSA v2: exact out-of-memory recovery for attention (engine, planner, devkit, native kernels)
 Requires-Python: >=3.9
 License-File: LICENSE

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/README.md

@@ -69,20 +69,20 @@ compiling: `CQSA_SKIP_EXT=1 pip install --no-build-isolation git+https://github.
 
 ### CUDA build (GitHub release page)
 
-The [release page](https://github.com/yiming-b/Stream-CQSA-v2/releases/tag/v2.2.2) carries, for
+The [release page](https://github.com/yiming-b/Stream-CQSA-v2/releases/tag/v2.2.3) carries, for
 every python 3.10/3.11/3.12 x torch 2.5.1/2.6.0 x CUDA 12.4 x GPU architecture (sm80 = A100 and
 other 8.x cards, sm90 = H100) cell, two wheels: `stream_cqsa` with the classic CUDA extensions and
 the companion `stream_cqsa_native` with the wave kernel. Install both; they replace the PyPI
 package with the same version plus the extensions:
 
 ```bash
-V=v2.2.2; T=1cu124torch2.6sm80; PY=cp311     # pick your torch / CUDA / GPU / python
-pip install https://github.com/yiming-b/Stream-CQSA-v2/releases/download/$V/stream_cqsa-2.2.2-$T-$PY-$PY-linux_x86_64.whl
-pip install https://github.com/yiming-b/Stream-CQSA-v2/releases/download/$V/stream_cqsa_native-2.2.2-$T-$PY-$PY-linux_x86_64.whl
+V=v2.2.3; T=1cu124torch2.6sm80; PY=cp311     # pick your torch / CUDA / GPU / python
+pip install https://github.com/yiming-b/Stream-CQSA-v2/releases/download/$V/stream_cqsa-2.2.3-$T-$PY-$PY-linux_x86_64.whl
+pip install https://github.com/yiming-b/Stream-CQSA-v2/releases/download/$V/stream_cqsa_native-2.2.3-$T-$PY-$PY-linux_x86_64.whl
 ```
 
 Also on the release page: one wheel with both extensions for torch 2.10 / CUDA 13 / python 3.11
-(`stream_cqsa-2.2.2-1cu130torch210sm8090-...whl`, sm80 + sm90), and the pure wheel and sdist that
+(`stream_cqsa-2.2.3-1cu130torch210sm8090-...whl`, sm80 + sm90), and the pure wheel and sdist that
 PyPI serves. The CUDA wheels are too large for PyPI (100 MB per file), which is why they live on
 GitHub, as PyTorch's own CUDA wheels do. They are built by `.github/workflows/wheels.yml` on every
 tag; publishing to PyPI is the `pypi` job, enabled by the repository variable `PYPI_PUBLISH=true`
@@ -94,49 +94,45 @@ fp16/bf16 and head dims 64/128); `CQSA_BUILD_NATIVE=1` adds the wave kernel.
 
 ## Use
 
+One function, with the signature of `torch.nn.functional.scaled_dot_product_attention`:
+
 ```python
 import torch, stream_cqsa
 
-q, k, v = (torch.randn(1, 8, 4_000_000, 64, dtype=torch.float16) for _ in range(3))   # host or device
-out = stream_cqsa.attention(q, k, v, is_causal=True)          # SDPA's signature; exact, always fits
-```
-
-That is the whole API for most uses. Below the memory boundary the call *is* the
-monolithic kernel; above it the planner chooses the decomposition from the free memory
-and the call runs exactly, on the best kernel available. Gradients flow when the inputs
-require them. To route an existing model without touching it:
-
-```python
-stream_cqsa.patch_sdpa()          # F.scaled_dot_product_attention now falls back to Stream-CQSA above 64K tokens
-with stream_cqsa.patched_sdpa():  # or scoped
-    model(...)
+q, k, v = (torch.randn(1, 8, 4_000_000, 64, dtype=torch.float16) for _ in range(3))   # host or device tensors
+out = stream_cqsa.attention(q, k, v, is_causal=True)          # exact, always fits
 ```
 
-Seeing what it does, and what it would cost:
+Below the memory boundary the call *is* the monolithic kernel; above it the planner chooses the
+decomposition from the free memory and the call runs exactly, on the best kernel installed, recovering
+from out-of-memory on its own. Gradients flow when the inputs require them. There is nothing else to learn;
+the rest is optional:
 
-```python
-out = stream_cqsa.attention(q, k, v, is_causal=True, verbose=True)   # or CQSA_VERBOSE=1
-#  Stream-CQSA: forward of N=4.0M tokens ... decomposed over c=31 at depth itr=1: 31 subproblems on cuda:0 | ... | expected ~2m10s
-#  Stream-CQSA: 100%|██████████| 31/31 [02:05<00:00,  4.05s/subproblem]
-#  Stream-CQSA: done in 2m06s (31 subproblems)
-stream_cqsa.estimate(16_777_216, H=8, D=64)     # dry run: monolithic fits?, chosen configuration, device/host memory, time
-stream_cqsa.calibrate()                         # ~1 min once per GPU model; saved to ~/.cache/stream_cqsa and used from then on
-```
-
-Explicit control (all optional keyword overrides of `attention`, or the functions
-themselves): `stream_cqsa_forward(q, k, v, itr=, c=, interest_set=, low_memory=,
-stream_from_host=, max_parallel=, verbose=)`, `stream_cqsa_backward(..., itr="auto")`,
-`stream_cqsa_attn` (autograd), `native_wave.wave_forward/wave_attention` (the wave kernel),
-`plan(...)`, `autotune(...)`, `devkit.compare_kernels(inner_fn, mono_fn)`,
-`adapters.flex_inner(score_mod, extra_mask_mod)`, `distributed.dist_stream_cqsa_forward/_backward`
-under `torch.distributed`. `attention(..., kernel="wave"|"cuda"|"triton")` pins the kernel.
+| I want to... | how |
+|---|---|
+| see what it decided and what ran | `attention(..., verbose=True)`, or `return_plan=True` / `return_info=True` |
+| know before running a huge call | `stream_cqsa.estimate(N=16_777_216)` (memory, configuration, time; nothing is allocated) |
+| be safe on a shared or small GPU | create inputs in host memory, or `q, k, v = stream_cqsa.place_inputs(q, k, v)` (moves what fits) |
+| run inside a model where Q/K/V fill the GPU | nothing: `attention` releases Q/K/V to pinned host memory when nothing fits next to them, puts the output on the GPU first, and brings Q/K/V back if room remains (`evict_inputs=`, `restore_inputs=` = `"auto"`｜`True`｜`False`; `del` a fused QKV tensor first so the memory can be freed) |
+| describe the budget instead of detecting it | `attention(..., hardware={"cuda:0": "10GiB", "host": "64GiB"})` |
+| route an existing model without touching it | `stream_cqsa.patch_sdpa()` or `with stream_cqsa.patched_sdpa(): model(...)` |
+| pin a setting | keywords of `attention`: `itr=`, `c=`/`interest_set=`, `accumulate_on_gpu=False`, `max_parallel=`, `bwd_itr=`, `kernel="cuda"｜"triton"｜"wave-cuda"｜"wave-triton"` |
+| make the planner's times accurate for my GPU | `stream_cqsa.calibrate()` once (about a minute; cached in `~/.cache/stream_cqsa`) |
+| check an installation | `stream-cqsa-doctor` |
+
+Everything else in the package is the engine layer that `attention` is built on, for kernel and systems
+developers: `stream_cqsa_forward` / `stream_cqsa_backward` / `stream_cqsa_attn` (the streaming engine),
+`native_wave` (the wave engine), `autoconfig.plan` / `autotune`, `devkit.compare_kernels(inner_fn, mono_fn)`,
+`adapters.flex_inner(score_mod, extra_mask_mod)`, and `distributed.dist_stream_cqsa_forward/_backward` under
+`torch.distributed`. `attention_oom_safe`, `stream_cqsa_auto` and `auto_attention` are legacy entry points kept
+for old code; they emit a `DeprecationWarning` pointing to `attention`.
 
 `sbatch slurm/quickstart.slurm` runs all of the above on one GPU (`benchmarks/quickstart.py`;
 output in `results/quickstart/`): the dry run, the monolithic path, the decomposed path under
 a memory cap with progress, host-resident inputs, autograd, and a patched module.
 
-Notebooks (executed, outputs included): `notebooks/stream_cqsa_v2_demo.ipynb` runs every
-feature on one GPU; `notebooks/oom_boundary_demo.ipynb` sweeps N explicitly under a memory cap
+Notebooks (executed, outputs included): `notebooks/stream_cqsa_v2_demo.ipynb` shows the one
+function (Part A) and the engine layer under it (Part B) on one GPU; `notebooks/oom_boundary_demo.ipynb` sweeps N explicitly under a memory cap
 and shows the baseline matching Stream-CQSA below the boundary and OOM-ing above it while
 Stream-CQSA continues, exact.
 

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa/__init__.py

@@ -1,4 +1,16 @@
-__version__ = "2.2.2"
+"""
+Stream-CQSA. The one function to use:
+
+    out = stream_cqsa.attention(q, k, v, is_causal=True)      # drop-in for F.scaled_dot_product_attention
+
+It runs the normal kernel when the call fits and the exact decomposition when it does not, plans the
+configuration, recovers from out-of-memory, supports autograd, and takes host or device inputs.
+Helpers around it: place_inputs, estimate, calibrate, patch_sdpa / patched_sdpa, doctor. Everything else
+in this package (stream_cqsa_forward, stream_cqsa_attn, native_wave, distributed, devkit, adapters) is the
+engine layer for kernel and systems developers; attention_oom_safe / stream_cqsa_auto / auto_attention
+are legacy entry points kept for old code.
+"""
+__version__ = "2.2.3"
 
 from .stable_stream import (
     stream_cqsa_forward,
@@ -10,6 +22,7 @@ from .oom_fallback import (stream_cqsa_auto, attention_oom_safe,
 from .native_autograd import stream_cqsa_attn, StreamCQSAAttention
 from .api import attention, estimate, patch_sdpa, unpatch_sdpa, patched_sdpa, kernels_available, place_inputs
 from .doctor import doctor
+from .autoconfig import calibrate
 from .progress import verbose_enabled
 
 from .interface import (

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa/api.py

@@ -94,6 +94,15 @@ def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask=None,
     "wave-cuda", "wave-triton", "cuda", "triton".
     ``overrides`` (itr=, c=, interest_set=, max_parallel=, ...) pin engine settings.
     ``plan_only=True`` returns the plan without computing anything.
+
+    Inside a model (Q/K/V already on the device, little room left): ``evict_inputs="auto"``
+    (default) releases Q/K/V to pinned host memory for the computation when nothing fits next
+    to them -- the caller's tensors are rebound, so the device memory is really freed -- and
+    streams from there. The output is placed on the device first; then ``restore_inputs="auto"``
+    brings Q/K/V back if room remains (otherwise they stay valid host tensors). ``True`` /
+    ``False`` force or forbid either step. Gradients still arrive on the device. Memory is only
+    freed if nothing else references the storage (e.g. ``del`` a fused QKV tensor first);
+    ``return_info=True`` reports ``inputs_evicted``, ``evict_released_gib``, ``inputs_restored``.
     """
     _INSIDE[0] += 1
     try:
@@ -159,6 +168,112 @@ def place_inputs(*tensors: torch.Tensor, device="cuda", reserve_fraction: float
     return tuple(out) if len(out) != 1 else out[0]
 
 
+def _pinned_copy(t: torch.Tensor) -> torch.Tensor:
+    """Host copy of a [B, H, N, D] device tensor, stored token-major ([B, N, H, D] in memory):
+    the layout the streaming engine gathers from, so it makes no second host copy."""
+    B, H, N, D = t.shape
+    try:
+        h = torch.empty((B, N, H, D), dtype=t.dtype, device="cpu", pin_memory=True)
+    except Exception:                                                # noqa: BLE001  (pinned memory exhausted / unavailable)
+        h = torch.empty((B, N, H, D), dtype=t.dtype, device="cpu")
+    h = h.transpose(1, 2)
+    h.copy_(t.detach())
+    return h
+
+
+class _ToHost(torch.autograd.Function):
+    """Differentiable move to pinned host memory; the gradient returns to the source device."""
+    @staticmethod
+    def forward(ctx, t):
+        ctx.dev = t.device
+        return _pinned_copy(t)
+
+    @staticmethod
+    def backward(ctx, g):
+        return g.to(ctx.dev)
+
+
+def _evict_inputs(tensors, grad: bool, vb: bool):
+    """Move the caller's device tensors to pinned host memory and RELEASE the device copies.
+
+    A plain ``.cpu()`` copies and leaves the device allocation alive behind the caller's own
+    reference. Here each tensor object is rebound onto the host copy (``t.data = host``), so the
+    device storage is freed unless something else still references it (a fused QKV projection
+    the slices are views of, an activation another op saved). What was actually released is
+    measured, not assumed. Returns (host_tensors, state); with ``grad`` the host tensors are
+    differentiable copies and the gradients come back on the device."""
+    dev = tensors[0].device
+    torch.cuda.synchronize(dev)
+    before = torch.cuda.memory_allocated(dev)
+    seen, expected = set(), 0
+    for t in tensors:
+        st = t.untyped_storage()
+        if st.data_ptr() not in seen:
+            seen.add(st.data_ptr()); expected += st.nbytes()
+    del st                                     # a storage handle keeps the device allocation alive
+    strides = [t.stride() for t in tensors]    # restored tensors get the caller's layout back
+    hosts = []
+    for t in tensors:
+        h = _ToHost.apply(t) if (grad and t.requires_grad) else _pinned_copy(t)
+        t.data = h.detach()                    # shares the host copy: no second host allocation
+        hosts.append(h)
+    torch.cuda.synchronize(dev)
+    torch.cuda.empty_cache()
+    freed = before - torch.cuda.memory_allocated(dev)
+    state = {"device": dev, "tensors": list(tensors), "strides": strides, "freed": int(freed), "expected": int(expected),
+             "restored": None}
+    if vb:
+        print(f"Stream-CQSA: evicted Q/K/V to pinned host memory, released {freed / 2**30:.2f} GiB on {dev}", flush=True)
+    if freed < 0.5 * expected:
+        warnings.warn(f"Stream-CQSA: evicting Q/K/V released only {freed / 2**30:.2f} of {expected / 2**30:.2f} GiB on {dev}: "
+                      "something else still references their storage (a fused QKV tensor they are views of, or an activation "
+                      "saved by another op). Drop that reference before the call (e.g. `del qkv`) to free the memory.",
+                      RuntimeWarning, stacklevel=4)
+    return hosts, state
+
+
+def _restore_inputs(state, mode, vb: bool, reserve_fraction: float = 0.10, budget_bytes: Optional[int] = None) -> bool:
+    """After the output is in place: bring Q/K/V back to the device if there is room (mode "auto"),
+    always try (True) or never (False). The output has priority, so this runs last. Tensors that
+    do not fit stay in host memory (``t.device`` tells the caller where each one is).
+
+    Room is judged on the memory this process may use, not the physical device: the free memory
+    under a set_per_process_memory_fraction cap, less ``reserve_fraction`` of the capped capacity,
+    and within ``budget_bytes`` (what a ``hardware=`` budget leaves after the output) when given.
+    A reserve of the physical total never restored anything under a small cap or budget."""
+    from .stable_stream import effective_free_bytes, effective_total_bytes
+    dev = state["device"]
+    if mode is False:
+        state["restored"] = False
+        return False
+    ok = True
+    left = budget_bytes
+    for t, stride in zip(state["tensors"], state["strides"]):
+        if t.device == dev:
+            continue
+        need = t.numel() * t.element_size()
+        room = effective_free_bytes(dev) - reserve_fraction * effective_total_bytes(dev)
+        if mode == "auto" and (need > room or (left is not None and need > left)):
+            ok = False
+            continue
+        try:
+            d = torch.empty_strided(t.shape, stride, dtype=t.dtype, device=dev)
+            d.copy_(t.data, non_blocking=True)
+            torch.cuda.synchronize(dev)
+            t.data = d
+            if left is not None:
+                left -= need
+        except Exception as exc:                                     # noqa: BLE001
+            if not _is_oom(exc):
+                raise
+            torch.cuda.empty_cache(); ok = False
+    state["restored"] = ok
+    if vb:
+        print("Stream-CQSA: Q/K/V restored to " + str(dev) if ok else
+              f"Stream-CQSA: Q/K/V left in host memory (no room on {dev} next to the output); they remain valid tensors", flush=True)
+    return ok
+
+
 def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, verbose, hardware, kernel, plan_only,
                return_plan, return_info=False, **overrides):
     """``return_info=True`` appends the engine's info dict (engine, kernel, itr, subproblem count,
@@ -176,11 +291,58 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
     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"
     out_el = q.element_size() if on_cuda else 0          # attention() delivers fp16 to a device caller, nothing to a host caller
+    # Inputs already on the device (the end-to-end model case). evict_inputs: "auto" moves them to
+    # host memory for the computation when nothing device-resident fits, True always, False never.
+    # restore_inputs: after the output is on the device, "auto" brings them back if there is room.
+    evict = overrides.pop("evict_inputs", "auto")
+    restore = overrides.pop("restore_inputs", "auto")
+    if evict not in ("auto", True, False) or restore not in ("auto", True, False):
+        raise ValueError("evict_inputs and restore_inputs take 'auto', True or False")
+    vb = verbose_enabled(verbose)
+    ev = {"state": None}                                  # eviction state, shared with _run (survives an exception there)
+    resident = (sum(t.numel() * t.element_size() for t in (q, k, v)) if on_cuda else 0)
+    pkw = dict(resident_input_bytes=resident, budget_excludes_inputs=hardware is None,
+               evict_inputs=bool(evict) and on_cuda)
+    if evict is True and on_cuda and not plan_only:
+        (q, k, v), ev["state"] = _evict_inputs((q, k, v), direction == "bwd", vb)
+        on_cuda = False
+        pkw = dict(resident_input_bytes=0)
+        if hardware is None:
+            hw = detect_hardware()
     p = _plan(N=N, B=B, H=H, D=D, dtype=q.dtype, causal=bool(is_causal), hardware=hw, direction=direction,
-              allow_distributed=False, out_bytes_per_el=out_el)
+              allow_distributed=False, out_bytes_per_el=out_el, **pkw)
     if plan_only:
         return p
-    vb = verbose_enabled(verbose)
+    out = None
+    try:
+        out, info = _run(q, k, v, p, is_causal=is_causal, scale=scale, direction=direction, on_cuda=on_cuda, out_el=out_el,
+                         overrides=overrides, kernel=kernel, verbose=verbose, vb=vb, evict=evict, ev=ev,
+                         out_device=out_device, pkw=pkw)
+    finally:
+        evicted = ev["state"]
+        if evicted is not None and evicted["restored"] is None:
+            # a hardware= budget describes the whole allowance: Q/K/V return within what the output leaves of it
+            budget = None if hardware is None else next(
+                (d.budget_bytes for d in hw.devices if torch.device(d.name) == evicted["device"]), None)
+            if budget is not None and out is not None and out.is_cuda:
+                budget -= out.numel() * out.element_size()
+            try:
+                _restore_inputs(evicted, restore, vb, budget_bytes=budget)   # the output is delivered first; Q/K/V only if room remains
+            except Exception:                                        # noqa: BLE001
+                pass
+    if evicted is not None:
+        info = dict(info, inputs_evicted=True, inputs_restored=bool(evicted["restored"]),
+                    evict_released_gib=round(evicted["freed"] / 2**30, 3))
+    else:
+        info = dict(info, inputs_evicted=False)
+    return _ret(out, p, info)
+
+
+def _run(q, k, v, p, *, is_causal, scale, direction, on_cuda, out_el, overrides, kernel, verbose, vb, evict, ev,
+         out_device, pkw):
+    """Execute the plan; returns (output on the caller's device, info)."""
+    from .autoconfig import plan as _plan, detect_hardware
+    B, H, N, D = q.shape
 
     # ---- monolithic: the call fits, so it IS the normal kernel ------------------
     if p.mode == "mono" and not overrides:
@@ -191,8 +353,8 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
         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 _ret(out, p, {"engine": "monolithic", "kernel": "sdpa", "itr": 0, "n_subproblems": 1})
+            out = out if out.device == out_device else out.to(out_device)     # evicted inputs: the caller's device
+            return out, {"engine": "monolithic", "kernel": "sdpa", "itr": 0, "n_subproblems": 1}
         except Exception as exc:                                     # noqa: BLE001
             if not _is_oom(exc):
                 raise
@@ -201,7 +363,8 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
             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, out_bytes_per_el=out_el)
+                      direction=direction, allow_distributed=False, out_bytes_per_el=out_el,
+                      **dict(pkw, budget_excludes_inputs=True))
             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)
 
@@ -217,6 +380,13 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
             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 evict:
+            # Nothing fits next to Q/K/V: release them to host memory for the computation. The
+            # output goes to the device first afterwards; Q/K/V return only if room remains.
+            (q, k, v), ev["state"] = _evict_inputs((q, k, v), direction == "bwd", vb)
+            on_cuda = False
+            if not kw.get("stream_from_host"):
+                kw.update(stream_from_host=True)
         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
@@ -231,16 +401,17 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
             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))
+                q, k, v = (_pinned_copy(t) for t in (q, k, v))            # token-major: the engine streams it as is
             on_cuda = False
     kw.update(overrides)
+    bwd_itr = kw.pop("bwd_itr", "auto")          # depth of the backward: "auto" (planned on its own), "fwd", or an int
     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)
         _acc = "host (CPU)" if (kw.get("low_memory") or kw.get("accumulate_on_gpu") is False) else "device (GPU)"
-        _qkv = "host, streamed" if kw.get("stream_from_host") else "device"
+        _qkv = ("device -> evicted to host, streamed" if ev["state"] else "host, streamed") if kw.get("stream_from_host") else "device"
         _itr = kw.get("itr", "auto"); _c = int(kw.get("c", 7))
         _nsub = (_c ** int(_itr)) if isinstance(_itr, int) else "planned per call"
         if use_wave:
@@ -256,7 +427,8 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
         wkern = WAVE_KERNELS.get(kernel, "auto")
         try:
             if direction == "bwd":
-                out = wave_attention(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, **wkw)
+                out = wave_attention(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose,
+                                     bwd_itr=(bwd_itr if isinstance(bwd_itr, int) else None), **wkw)
                 info = {"engine": "wave", "kernel": wkern, **wkw}
             else:
                 acc_gpu = not (kw.get("low_memory") or kw.get("accumulate_on_gpu") is False)
@@ -282,7 +454,7 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
                                    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"))
+                                   max_parallel=kw.get("max_parallel"), bwd_itr=bwd_itr)
             info = {"engine": "classic", "kernel": "triton" if os.environ.get("CQSA_FORWARD", "").lower() == "triton" else "cuda",
                     "itr": kw.get("itr", "auto"), "c": kw.get("c", 7), "stream_from_host": bool(kw.get("stream_from_host", host)),
                     "accumulate_on_gpu": bool(kw.get("accumulate_on_gpu", True))}
@@ -294,13 +466,14 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
             info["engine"] = "classic"; info["kernel"] = "triton" if os.environ.get("CQSA_FORWARD", "").lower() == "triton" else "cuda"
     if vb:
         print(_config_line(stage="ran", mode="cqsa", itr=info.get("itr", kw.get("itr")), c=info.get("c", kw.get("c", 7)),
-                           n_sub=info.get("n_subproblems"), qkv=("host, streamed" if kw.get("stream_from_host") else "device"),
+                           n_sub=info.get("n_subproblems"),
+                           qkv=(("device -> evicted to host, streamed" if ev["state"] else "host, streamed") if kw.get("stream_from_host") else "device"),
                            acc=("host (CPU)" if (kw.get("low_memory") or kw.get("accumulate_on_gpu") is False) else "device (GPU)"),
                            engine=info.get("engine", "?"), kernel=info.get("kernel", "?"), n_par=info.get("n_parallel", kw.get("max_parallel")),
                            out=f"{out_device} {str(q.dtype).replace('torch.', '')}")
               + (f" | OOM retries {info['oom_retries']}" if info.get("oom_retries") else ""), flush=True)
     out = _deliver(out, out_device, q.dtype, vb)
-    return _ret(out, p, info)
+    return out, info
 
 
 def _deliver(out: torch.Tensor, device, dtype, vb: bool) -> torch.Tensor:

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa/autoconfig.py

@@ -321,11 +321,12 @@ class Plan:
     est_host_gib: float = 0.0
     reason: str = ""
     candidates: list[dict] = field(default_factory=list)
+    evict_inputs: bool = False     # Q/K/V are on the device now and are moved to host memory for the computation
 
     def name(self) -> str:
         if self.mode == "mono":
             return "monolithic"
-        s = f"cqsa c={self.c} itr={self.itr} acc={self.acc} n_par={self.n_par}" + (" host-resident Q/K/V" if self.stream_from_host else "")
+        s = f"cqsa c={self.c} itr={self.itr} acc={self.acc} n_par={self.n_par}" + ((" Q/K/V evicted to host" if self.evict_inputs else " host-resident Q/K/V") if self.stream_from_host else "")
         return s + (f" on {self.world} devices" if self.world > 1 else "")
 
     def engine_kwargs(self) -> dict:
@@ -348,7 +349,8 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
          hardware: HardwareSpec | None = None, model: CostModel | None = None,
          max_itr: int = 3, n_pars: Sequence[int] = (1, 2, 4), allow_distributed: bool = True,
          quorum_sets: dict[int, Sequence[int]] | None = None, direction: str = "fwd",
-         safety: float = 0.85, tol: float = 0.02, out_bytes_per_el: int = 4) -> Plan:
+         safety: float = 0.85, tol: float = 0.02, out_bytes_per_el: int = 4,
+         resident_input_bytes: int = 0, budget_excludes_inputs: bool = True, evict_inputs: bool = False) -> Plan:
     """
     Choose the configuration for one attention call on `hardware` (detected if None).
     ``out_bytes_per_el``: device bytes per element the delivered forward output costs when the
@@ -362,6 +364,14 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
     (l^2/c)^itr and per-task overhead is c^itr; the model has both terms.
     `direction="bwd"` plans the backward (its own memory estimate and cost) --
     the forward and backward depths are chosen independently.
+
+    Inputs already on the device (the end-to-end model case): ``resident_input_bytes`` is what
+    Q/K/V occupy there now. ``budget_excludes_inputs`` says whether the budget was measured with
+    them in place (free memory: True) or describes the whole allowance (a hardware dict: False).
+    A device-resident candidate counts Q/K/V in its estimate, so when the budget excludes them
+    they are credited back. A host-streaming candidate needs them off the device: with
+    ``evict_inputs`` their bytes are freed for the computation (the candidate is marked
+    ``evict``); without it they stay and are charged on top of the estimate.
     """
     hw = hardware or detect_hardware()
     cm = model or load_cost_model() or CostModel()
@@ -375,6 +385,8 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
         groups.setdefault(d.identical_group, []).append(d)
     devs = max(groups.values(), key=lambda g: (len(g), min(x.budget_bytes for x in g)))
     dev_budget = min(d.budget_bytes for d in devs) * safety
+    R = int(resident_input_bytes)
+    dev_budget += R if budget_excludes_inputs else 0          # estimates below count Q/K/V; the measured budget did not
     cands: list[dict] = []
 
     bwd = direction == "bwd"
@@ -406,19 +418,26 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
                                        stream_from_host=host, accumulate_on_gpu=(acc == "gpu"), **extra)
                             hostb = _host_bytes(N, B, H, D, itemsize, host, acc) * (2.0 if bwd else 1.0)
                             t = cm.cqsa_time(N, B, H, D, itr, causal, acc, host, n_par, world=w, hw=hw, c=c, l=l, direction=direction)
+                            evict = bool(host and R and evict_inputs)
+                            stay = R if (host and R and not evict) else 0      # streaming next to the unevicted device copy
                             cands.append(dict(mode="cqsa" if w == 1 else "cqsa_dist", itr=itr, c=c, interest_set=tuple(iset),
-                                              acc=acc, stream_from_host=host,
-                                              n_par=n_par, world=w, time=t, peak=peak / GIB, host=hostb / GIB,
-                                              ok=(peak <= dev_budget and hostb <= hw.host_budget_bytes)))
+                                              acc=acc, stream_from_host=host, evict=evict,
+                                              n_par=n_par, world=w, time=t, peak=(peak + stay) / GIB, host=hostb / GIB,
+                                              ok=(peak + stay <= dev_budget and hostb <= hw.host_budget_bytes)))
     feas = [c for c in cands if c["ok"]]
     if not feas:
         floor = per_el_floor(itemsize, True, False) * N * B * H * D
-        p = Plan(mode="cqsa", itr=max_itr, c=max(qsets), interest_set=tuple(qsets[max(qsets)]), direction=direction,
-                 acc="cpu", stream_from_host=True, n_par=1, world=len(devs),
-                 devices=[d.name for d in devs], candidates=cands,
+        # the most frugal configuration that was enumerated (so the subproblem-count cap holds): everything
+        # offloaded, one subproblem in flight, smallest predicted peak
+        off = [c for c in cands if c["mode"] == "cqsa" and c["stream_from_host"] and c["acc"] == "cpu" and c["n_par"] == 1]
+        fb = min(off, key=lambda c: (c["peak"], c["time"]))
+        p = Plan(mode="cqsa", itr=fb["itr"], c=fb["c"], interest_set=tuple(fb["interest_set"]), direction=direction,
+                 acc="cpu", stream_from_host=True, n_par=1, world=1, evict_inputs=bool(R and evict_inputs),
+                 devices=[d.name for d in devs[:1]], candidates=cands,
                  reason=(f"nothing fits: device budget {dev_budget / GIB:.1f} GiB, host budget {hw.host_budget_bytes / GIB:.0f} GiB; "
                          f"the O(N.H.D) device floor with everything offloaded is {floor / GIB:.1f} GiB. "
-                         f"Falling back to the deepest, most offloaded configuration on every device; expect OOM."))
+                         f"Falling back to the most frugal configuration (c={fb['c']}, itr={fb['itr']}, everything offloaded, "
+                         f"~{fb['peak']:.1f} GiB); the engine escalates on its own if that still does not fit."))
         return p
     front = [c for c in feas if not any((o["time"] < c["time"]) and (o["peak"] < c["peak"]) for o in feas)]
     best_t = min(c["time"] for c in front)
@@ -427,7 +446,8 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
     p = Plan(mode=pick["mode"], itr=pick["itr"], c=pick["c"], interest_set=tuple(pick["interest_set"]), direction=direction,
              acc=pick["acc"], stream_from_host=pick["stream_from_host"],
              n_par=pick["n_par"], world=pick["world"], devices=[d.name for d in devs[:pick["world"]]],
-             est_time_s=pick["time"], est_peak_gib=pick["peak"], est_host_gib=pick["host"], candidates=cands)
+             est_time_s=pick["time"], est_peak_gib=pick["peak"], est_host_gib=pick["host"], candidates=cands,
+             evict_inputs=bool(pick.get("evict", False)))
     mono_c = cands[0]
     if pick["mode"] == "mono":
         p.reason = (f"monolithic fits ({mono_c['peak']:.2f} GiB <= {dev_budget / GIB:.1f} GiB budget) and is the fastest "
@@ -572,6 +592,9 @@ def auto_attention(q, k, v, *, causal: bool = False, scale: float | None = None,
     the planned size; otherwise the single-device frontier point is used.
     `overrides` are forwarded to the engine (e.g. allow_escalation=False).
     """
+    import warnings
+    warnings.warn("stream_cqsa.autoconfig.auto_attention is a legacy entry point; use stream_cqsa.attention(q, k, v, "
+                  "is_causal=..., hardware=..., return_plan=True)", DeprecationWarning, stacklevel=2)
     from .stable_stream import stream_cqsa_forward
     if isinstance(hardware, dict):
         hardware = hardware_from_dict(hardware)

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa/oom_fallback.py

@@ -87,6 +87,13 @@ def oom_hint(N: int, B: int, H: int, D: int, itemsize: int, device=None) -> str:
     return "; ".join(parts)
 
 
+def _deprecated_entry(name: str):
+    import warnings
+    warnings.warn(f"stream_cqsa.{name} is a legacy entry point; use stream_cqsa.attention(q, k, v, is_causal=...), "
+                  "which plans, falls back on out-of-memory, supports autograd and host or device inputs.",
+                  DeprecationWarning, stacklevel=3)
+
+
 def attention_oom_safe(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -110,6 +117,7 @@ def attention_oom_safe(
     permits moving Q/K/V to host memory on the fallback -- see the module
     docstring; without it, inputs that do not fit cannot be recovered.
     """
+    _deprecated_entry("attention_oom_safe")
     from .stable_stream import max_depth_for, stream_cqsa_forward
 
     if q.dim() != 4:
@@ -312,6 +320,7 @@ def stream_cqsa_auto(q, k, v, *, causal=False, scale=None, return_info=False,
     `return_info=True`; `info["lse"]` is required by `stream_cqsa_backward`,
     and `info["config"]` reports which rung ran).
     """
+    _deprecated_entry("stream_cqsa_auto")
     from .stable_stream import stream_cqsa_forward
 
     out_dtype = q.dtype

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa/stable_stream.py

@@ -644,6 +644,19 @@ def estimate_peak_bytes(N: int, itr: int, *, B: int, H: int, D: int, itemsize: i
     return int(floor + transient)
 
 
+def effective_total_bytes(device) -> int:
+    """Device capacity as the allocator will grant it: the physical total times the
+    torch.cuda.set_per_process_memory_fraction cap (the whole device when none is set)."""
+    idx = device.index if isinstance(device, torch.device) else device
+    if idx is None or isinstance(idx, str):
+        idx = torch.cuda.current_device()
+    try:
+        frac = torch.cuda.get_per_process_memory_fraction(int(idx))
+    except Exception:
+        frac = 1.0
+    return int(min(frac, 1.0) * torch.cuda.get_device_properties(int(idx)).total_memory)
+
+
 def effective_free_bytes(device) -> int:
     """
     Free device memory as the allocator will actually grant it: the driver's

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3/stream_cqsa.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stream-cqsa
-Version: 2.2.2
+Version: 2.2.3
 Summary: Stream-CQSA v2: exact out-of-memory recovery for attention (engine, planner, devkit, native kernels)
 Requires-Python: >=3.9
 License-File: LICENSE

{stream_cqsa-2.2.2 → stream_cqsa-2.2.3}/stream_cqsa.egg-info/SOURCES.txt

@@ -1897,6 +1897,7 @@ tests/test_bf16.py
 tests/test_bwd_host_stream.py
 tests/test_devkit_autoconfig.py
 tests/test_escalation_autograd.py
+tests/test_evict_inputs.py
 tests/test_place_and_verbose.py
 tests/test_rect_ooc.py
 tests/test_reference.py