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

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

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: stream-cqsa
-Version: 2.2.1
+Version: 2.2.2
 Summary: Stream-CQSA v2: exact out-of-memory recovery for attention (engine, planner, devkit, native kernels)
 Requires-Python: >=3.9
 License-File: LICENSE
 Requires-Dist: torch
 Requires-Dist: numpy
+Requires-Dist: triton>=3.0; sys_platform == "linux"
+Requires-Dist: triton-windows>=3.0; sys_platform == "win32"
 Provides-Extra: triton
 Requires-Dist: triton; extra == "triton"
 Provides-Extra: progress

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

@@ -46,7 +46,8 @@ Two implementations of the kernels ship, and where you get them differs:
 | **CUDA build** | the classic CUDA extensions (FlashAttention-2-derived CQS forward, CUDA backward) plus the native wave kernel; the fastest forward | GitHub release page (prebuilt wheels per python / torch / CUDA / GPU) or a source build | see below |
 
 The package picks at run time: CUDA extensions and the wave kernel when they are importable,
-Triton otherwise. `stream-cqsa-doctor` shows which kernels loaded. Since v2.2.0 the wave engine
+Triton otherwise. `stream-cqsa-doctor` shows which kernels loaded. `stream_cqsa.place_inputs(q, k, v)` moves host tensors to the GPU only if they fit and keeps the rest pinned
+in host memory (a guardrail for shared GPUs); `verbose=True` prints every automatic choice. Since v2.2.0 the wave engine
 also runs on the Triton kernel (`attention(..., kernel="wave-triton")`; `"wave-cuda"` pins the
 CUDA one), supports the host accumulator (`accumulate_on_gpu=False`) on both, and the quorum-set
 register includes c=3 and c=133; the automatic planner enumerates every registered set. The Triton backward is the
@@ -61,26 +62,27 @@ pip install stream-cqsa                                                  # Trito
 python -m stream_cqsa.doctor                                             # what this machine can run, and how far
 ```
 
-`pip install triton` if your torch did not bring it. `pip install flash-attn` is optional and
+Triton is installed with the package (`triton` on Linux, `triton-windows` on Windows, where the official
+package has no wheels). `pip install flash-attn` is optional and
 gives the monolithic fast path its own kernel. The same build installs from a checkout without
 compiling: `CQSA_SKIP_EXT=1 pip install --no-build-isolation git+https://github.com/yiming-b/Stream-CQSA-v2.git`.
 
 ### CUDA build (GitHub release page)
 
-The [release page](https://github.com/yiming-b/Stream-CQSA-v2/releases/tag/v2.2.1) carries, for
+The [release page](https://github.com/yiming-b/Stream-CQSA-v2/releases/tag/v2.2.2) 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.1; 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.1-$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.1-$T-$PY-$PY-linux_x86_64.whl
+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
 ```
 
 Also on the release page: one wheel with both extensions for torch 2.10 / CUDA 13 / python 3.11
-(`stream_cqsa-2.2.1-1cu130torch210sm8090-...whl`, sm80 + sm90), and the pure wheel and sdist that
+(`stream_cqsa-2.2.2-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`
@@ -141,7 +143,7 @@ Stream-CQSA continues, exact.
 ## Layout
 
     stream_cqsa/      the package (engine, planner, devkit, adapters, distributed, autograd)
-    csrc/, csrc_nc/   the two kernel source trees (FlashAttention-2 + CQS; v11 and v9)
+    csrc/, csrc_nc/   the classic CUDA kernel's causal and non-causal builds (FlashAttention-2 + CQS)
     native/           the multi-subproblem "wave" kernel (cqsa_native): sources, build, tests, benchmark
     benchmarks/profile_sweep.py   one-parameter profiles (c; N x {itr, acc}) with quadratic/linear fits: results/profile_sweep/
     tests/            pytest suite (engine, planner, devkit, adapters)

{stream_cqsa-2.2.1 → stream_cqsa-2.2.2}/setup.py

@@ -247,9 +247,13 @@ setup(
            if os.getenv("CQSA_BUILD_NATIVE", "").strip().lower() in ("1", "true", "yes") else [])),
     cmdclass={} if SKIP_EXT else {"build_ext": BuildExtension.with_options(use_ninja=True)},
     python_requires=">=3.9",
-    install_requires=["torch", "numpy"],
+    # Triton is the no-build kernel path. PyTorch's Linux wheels already depend on it; on Windows the
+    # official package has no wheels, so the community build `triton-windows` is used instead.
+    install_requires=["torch", "numpy",
+                      'triton>=3.0; sys_platform == "linux"',
+                      'triton-windows>=3.0; sys_platform == "win32"'],
     extras_require={
-        "triton": ["triton"],                      # the no-build kernels
+        "triton": ["triton"],                      # the no-build kernels (explicit, any platform)
         "progress": ["tqdm"],                      # verbose=True progress bars (a plain bar is used without it)
         "bench": ["matplotlib", "tqdm"],
         "test": ["pytest"],

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

@@ -1,4 +1,4 @@
-__version__ = "2.2.1"
+__version__ = "2.2.2"
 
 from .stable_stream import (
     stream_cqsa_forward,
@@ -8,7 +8,7 @@ from .stable_stream import (
 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 .api import attention, estimate, patch_sdpa, unpatch_sdpa, patched_sdpa, kernels_available, place_inputs
 from .doctor import doctor
 from .progress import verbose_enabled
 

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

@@ -82,7 +82,7 @@ def _is_oom(exc: BaseException) -> bool:
 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:
+              return_plan: bool = False, return_info: 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).
@@ -98,13 +98,74 @@ def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask=None,
     _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)
+                          kernel=kernel, plan_only=plan_only, return_plan=return_plan, return_info=return_info, **overrides)
     finally:
         _INSIDE[0] -= 1
 
 
+def _config_line(*, stage: str, mode: str, itr=None, c=None, n_sub=None, qkv: str, acc: str, engine: str, kernel: str,
+                 n_par=None, out: str) -> str:
+    """One uniform line that states every automatic choice (printed with verbose=True)."""
+    dec = "monolithic (one kernel call, no decomposition)" if mode == "mono" else f"decomposed: itr={itr}, c={c}, {n_sub} subproblems"
+    parts = [f"Stream-CQSA {stage}: {dec}", f"Q/K/V on {qkv}", f"accumulator on {acc}", f"engine {engine}, kernel {kernel}"]
+    if n_par:
+        parts.append(f"{n_par} in flight")
+    parts.append(f"output to {out}")
+    return " | ".join(parts)
+
+
+def place_inputs(*tensors: torch.Tensor, device="cuda", reserve_fraction: float = 0.15, pin: bool = True,
+                 verbose: Optional[bool] = None):
+    """
+    Guardrail for a crowded device: move host tensors to ``device`` only if they fit, and keep
+    the rest in (pinned) host memory, where ``attention`` streams them from. A tensor is moved
+    when the free device memory minus ``reserve_fraction`` of the total still holds it; an
+    out-of-memory during the move is caught and the tensor stays on the host. Device tensors
+    are returned unchanged. Returns the tensors in the order given.
+
+        q, k, v = stream_cqsa.place_inputs(q, k, v)          # each on the GPU if it fits, else host (pinned)
+        out = stream_cqsa.attention(q, k, v, is_causal=True)
+    """
+    vb = verbose_enabled(verbose)
+    dev = torch.device(device)
+    out = []
+    for i, t in enumerate(tensors):
+        if t.is_cuda or dev.type != "cuda" or not torch.cuda.is_available():
+            out.append(t); continue
+        free, total = torch.cuda.mem_get_info(dev)
+        need = t.numel() * t.element_size()
+        moved = None
+        if need <= free - reserve_fraction * total:
+            try:
+                moved = t.to(dev)
+            except Exception as exc:                                     # noqa: BLE001
+                if not _is_oom(exc):
+                    raise
+                torch.cuda.empty_cache()
+        if moved is None:
+            if pin and not t.is_pinned():
+                try:
+                    t = t.pin_memory()
+                except Exception:                                        # noqa: BLE001
+                    pass
+            if vb:
+                print(f"Stream-CQSA place_inputs: tensor {i} ({need / 2**30:.2f} GiB) stays in host memory "
+                      f"({free / 2**30:.2f} GiB free on {dev}); attention() will stream it", flush=True)
+            out.append(t)
+        else:
+            if vb:
+                print(f"Stream-CQSA place_inputs: tensor {i} ({need / 2**30:.2f} GiB) -> {dev}", flush=True)
+            out.append(moved)
+    return tuple(out) if len(out) != 1 else out[0]
+
+
 def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, verbose, hardware, kernel, plan_only,
-               return_plan, **overrides):
+               return_plan, return_info=False, **overrides):
+    """``return_info=True`` appends the engine's info dict (engine, kernel, itr, subproblem count,
+    OOM retries, ...) to the return value: (out[, plan], info)."""
+    def _ret(out, p, info):
+        r = (out, p) if return_plan else out
+        return ((r + (info,)) if return_plan else (out, info)) if return_info else r
     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
@@ -114,8 +175,9 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
     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"
+    out_el = q.element_size() if on_cuda else 0          # attention() delivers fp16 to a device caller, nothing to a host caller
     p = _plan(N=N, B=B, H=H, D=D, dtype=q.dtype, causal=bool(is_causal), hardware=hw, direction=direction,
-              allow_distributed=False)
+              allow_distributed=False, out_bytes_per_el=out_el)
     if plan_only:
         return p
     vb = verbose_enabled(verbose)
@@ -123,12 +185,14 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
     # ---- 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)
+            print(f"Stream-CQSA: N={N} fits ({p.est_peak_gib:.1f} GiB)", flush=True)
+            print(_config_line(stage="plan", mode="mono", qkv=("device" if on_cuda else "host, copied to the device"),
+                               acc="device", engine="monolithic", kernel="sdpa", out=f"{q.device} {str(q.dtype).replace('torch.', '')}"), 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
+            return _ret(out, p, {"engine": "monolithic", "kernel": "sdpa", "itr": 0, "n_subproblems": 1})
         except Exception as exc:                                     # noqa: BLE001
             if not _is_oom(exc):
                 raise
@@ -137,7 +201,7 @@ 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)
+                      direction=direction, allow_distributed=False, out_bytes_per_el=out_el)
             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)
 
@@ -175,6 +239,17 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
     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"
+        _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:
+            from .native_wave import resolve_wave_kernel
+            _eng, _ker = "wave", resolve_wave_kernel(WAVE_KERNELS.get(kernel, "auto"), q.dtype, int(q.shape[-1]))
+        else:
+            _eng, _ker = "classic", ("triton" if (kernel in ("triton", "wave-triton") or os.environ.get("CQSA_FORWARD", "").lower() == "triton") else "cuda")
+        print(_config_line(stage="plan", mode="cqsa", itr=_itr, c=_c, n_sub=_nsub, qkv=_qkv, acc=_acc, engine=_eng, kernel=_ker,
+                           n_par=kw.get("max_parallel"), out=f"{out_device} {str(q.dtype).replace('torch.', '')}"), 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))))
@@ -182,10 +257,12 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
         try:
             if direction == "bwd":
                 out = wave_attention(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, **wkw)
+                info = {"engine": "wave", "kernel": wkern, **wkw}
             else:
                 acc_gpu = not (kw.get("low_memory") or kw.get("accumulate_on_gpu") is False)
-                out, _ = wave_forward(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose,
-                                      accumulate_on_gpu=acc_gpu, kernel=wkern, **wkw)
+                out, info = wave_forward(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose,
+                                         accumulate_on_gpu=acc_gpu, kernel=wkern, **wkw)
+                info = {k_: v_ for k_, v_ in info.items() if k_ != "lse"}; info["engine"] = "wave"
         except Exception as exc:                                     # noqa: BLE001
             if not _is_oom(exc) or kernel in WAVE_KERNELS:
                 raise
@@ -206,12 +283,24 @@ def _attention(q, k, v, attn_mask, dropout_p, is_causal, scale, enable_gqa, *, v
                                    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"))
+            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))}
         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, info = stream_cqsa_forward(q, k, v, causal=bool(is_causal), scale=scale, verbose=verbose, out_device="acc", **kw)
+            info = {k_: v_ for k_, v_ in info.items() if k_ not in ("lse",) and not hasattr(v_, "device")}
+            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"),
+                           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 (out, p) if return_plan else out
+    return _ret(out, p, info)
 
 
 def _deliver(out: torch.Tensor, device, dtype, vb: bool) -> torch.Tensor:

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

@@ -348,9 +348,11 @@ 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) -> Plan:
+         safety: float = 0.85, tol: float = 0.02, out_bytes_per_el: int = 4) -> 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
+    accumulator is on the host (4 = classic engine default, 2 = fp16 to a device caller, 0 = host caller).
 
     `quorum_sets` ({c: interest_set}, default QUORUM_SETS) is the fourth axis:
     every (c, itr) pair is a candidate. Measured at N=131K (A100-40GB): for the
@@ -399,8 +401,9 @@ def plan(*, N: int, B: int = 1, H: int = 8, D: int = 64, dtype=torch.float16, ca
                             if n_par > n_tasks:
                                 continue
                             est = estimate_peak_bytes_bwd if bwd else estimate_peak_bytes
+                            extra = {} if bwd else {"out_bytes_per_el": int(out_bytes_per_el)}
                             peak = est(N, itr, B=B, H=H, D=D, itemsize=itemsize, n_par=n_par, c=c, l=l,
-                                       stream_from_host=host, accumulate_on_gpu=(acc == "gpu"))
+                                       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)
                             cands.append(dict(mode="cqsa" if w == 1 else "cqsa_dist", itr=itr, c=c, interest_set=tuple(iset),