Publish Hydra kernel source packet

CARD.md

@@ -0,0 +1,67 @@
+# Hydra Kernel Card
+
+## Summary
+
+Hydra provides a bounded-residency decode attention path for long-context
+inference. The implementation is Python plus Triton and is packaged here as a
+Hugging Face universal CUDA kernel source directory.
+
+Source code: [https://github.com/newjordan/hydra/tree/main/hf-kernels/hydra](https://github.com/newjordan/hydra/tree/main/hf-kernels/hydra)
+
+## Intended Use
+
+Use Hydra for experiments where full decode attention over a long KV cache is
+memory-bound and a bounded resident set is acceptable for evaluation.
+
+Hydra is not intended as a drop-in universal FlashAttention replacement. Users
+should keep an exact-model fallback path and validate quality for their prompt
+set.
+
+## Kernel Interface
+
+The exported call is:
+
+```python
+hydra.hydra(q, k, v, is_causal=True, sliding_window=None)
+```
+
+Inputs are bf16 tensors shaped `(B, H, T, D)` with `D=128`. The decode path
+supports `Tq == 1`; the prefill path requires sequence length to be a multiple
+of the compile-time block size.
+
+## Evidence
+
+This card separates the kernel contribution from benchmark appendices:
+
+- kernel/package validation: import, CSR, CUDA decode parity, builder, example,
+  and isolated decode benchmark gates
+- broad Hydra campaign context: multi-GPU bounded-residency testing, comparison
+  lanes, capacity/OOM boundaries, and diagnostics in the staging repo
+- exact-Qwen proof-of-concept: summary-backed demo rows for:
+
+- RTX PRO 6000 WS with `Qwen/Qwen3.6-35B-A3B-FP8`
+- RTX 3090 with `Qwen/Qwen3.6-35B-A3B-FP8`
+
+Use `results/reports/QWEN3P6_FP8_EVIDENCE_TABLE.md` in the staging repo as the
+claim ledger for the exact-Qwen proof-of-concept only. The table is generated
+from raw summary JSON, answer artifacts, and logs. It intentionally excludes
+incomplete scopes from completed benchmark rows.
+
+## Non-Claims
+
+- no universal speedup claim
+- no production-readiness claim
+- no broad quality-preservation claim without scorer or inspection evidence
+- no proxy/profile/loader-only benchmark claims
+- no results from non-Qwen or non-FP8 runs in the exact-Qwen proof-of-concept table
+- no framing that treats the exact-Qwen proof-of-concept as the full Hydra campaign
+
+## Required Validation
+
+For source changes, run:
+
+- import and CSR tests
+- CUDA decode parity against PyTorch SDPA on small tensors
+- kernel-builder `ci-test`
+- one isolated decode benchmark
+- one exact-model reproduction on a named GPU/config

README.md

@@ -0,0 +1,158 @@
+---
+license: mit
+tags:
+  - kernel
+  - kernels
+  - triton
+  - attention
+  - long-context
+---
+
+# Hydra
+
+Hydra is an experimental bounded-residency attention kernel for long-context
+decode. It keeps sink tokens, recent tokens, and selected older pages resident
+instead of forcing each decode step to attend over the full KV cache.
+
+Source code: [https://github.com/newjordan/hydra/tree/main/hf-kernels/hydra](https://github.com/newjordan/hydra/tree/main/hf-kernels/hydra)
+
+This release is intentionally narrow. It is not a general replacement for
+full attention, and it does not claim universal speedups or broad quality
+preservation. The current target is fit and usability for specific
+long-context inference workloads where the full-attention path is memory-bound.
+
+## Usage
+
+After the kernel is published:
+
+```python
+import torch
+from kernels import get_kernel
+
+hydra = get_kernel("Frosty40/hydra")
+
+q = torch.randn(1, 32, 1, 128, device="cuda", dtype=torch.bfloat16)
+k = torch.randn(1, 8, 8192, 128, device="cuda", dtype=torch.bfloat16)
+v = torch.randn(1, 8, 8192, 128, device="cuda", dtype=torch.bfloat16)
+
+out = hydra.hydra(q, k, v)
+print(out.shape)
+```
+
+For local development from the public source checkout:
+
+```python
+from pathlib import Path
+import sys
+
+sys.path.insert(0, str(Path("hf-kernels") / "hydra" / "torch-ext"))
+import hydra
+```
+
+`readme_example.py` uses the local source packet by default so it can run before
+publication. Set `HYDRA_USE_HUB=1` after publication to exercise the Hub-loaded
+path.
+
+## API
+
+```python
+hydra.hydra(
+    q,
+    k,
+    v,
+    *,
+    is_causal=True,
+    sliding_window=None,
+    policy_layer_idx=None,
+    precision="high",
+)
+```
+
+Current constraints:
+
+- CUDA tensors only
+- bf16 `q`, `k`, and `v`
+- shape `(B, H, T, D)` with `D=128`
+- causal attention only
+- decode path supports `Tq == 1` with arbitrary `Tkv`
+- prefill path requires `T % BLOCK_SIZE == 0`
+
+## Evidence Boundary
+
+Submission-facing evidence must come from checked artifacts, not prose notes.
+Treat evidence in three separate scopes:
+
+- kernel/package validation: tests, CUDA parity logs, `kernel-builder` logs, and
+  isolated decode benchmarks for this source packet
+- broad Hydra research campaign: capacity, quality, sparse-attention comparison,
+  edge/OOM, diagnostic, and model-family reports from the staging repo
+- exact-model proof-of-concept: checked `Qwen/Qwen3.6-35B-A3B-FP8` rows for
+  named GPUs only
+
+The exact-Qwen proof-of-concept appendix in the staging repo is under:
+
+```text
+results/raw/qwen3p6_35b_a3b_fp8/
+results/reports/QWEN3P6_FP8_EVIDENCE_TABLE.md
+```
+
+Each cited row must include all three:
+
+- fit/headroom: GPU, context length, memory allocated/reserved, and OOM state
+- quality/correctness: prompt/task ID and generated answer artifact
+- speed/usability: wall time, generated tokens, tokens/sec, and comparison target
+
+Do not cite proxy models, loader-only probes, failed dependency checks, or
+non-matching model runs as Hydra benchmark results. Do not describe the
+exact-Qwen proof-of-concept subset as the full Hydra validation campaign.
+
+## Current Proof-Of-Concept Scope
+
+The current exact-Qwen artifact-backed proof-of-concept scope is:
+
+| GPU | Model | Scope |
+| --- | --- | --- |
+| RTX PRO 6000 WS | `Qwen/Qwen3.6-35B-A3B-FP8` | 32k/80k/160k repeat packet, 160k c96 warm packet, and frontier/headroom sweeps |
+| RTX 3090 | `Qwen/Qwen3.6-35B-A3B-FP8` | 2k/3k/4k/6k/8k fit probes and completed 10k/12k/14k edge sweep |
+
+The 3090 result should be framed as fit/usability evidence, not a speedup
+claim. Token rates are slow in the long-context edge rows. The broader Hydra
+campaign includes additional GPUs, tasks, and comparison lanes outside this
+exact-model appendix.
+
+## Validation Required Before Merge
+
+Minimum gates for source changes:
+
+```bash
+cd hf-kernels/hydra
+python3 -m pytest -q tests
+nix run .#ci-test
+python3 benchmarks/benchmark_hydra_decode.py --repo .
+python3 readme_example.py
+```
+
+Run the CUDA tests on real GPUs. Local syntax checks are not enough for a
+kernel submission.
+
+## Benchmark Snapshot
+
+The current 8192-token decode smoke/benchmark matrix is intentionally reported
+as kernel/package evidence, not as a universal speedup claim.
+
+| GPU | Package smoke decode | HF benchmark mean |
+| --- | ---: | ---: |
+| RTX 3060 | 0.2574 ms | 0.3229 ms |
+| RTX 3070 | 0.1474 ms | 0.2532 ms |
+| RTX 3080 | 0.2051 ms | 0.3157 ms |
+| RTX 3090 | 0.1492 ms | 0.3107 ms |
+| RTX 4070 Ti | 0.1261 ms | 0.2215 ms |
+| RTX 4090 | 0.1132 ms | 0.2245 ms |
+| A100 SXM4 | 0.1408 ms | 0.2568 ms |
+| RTX PRO 6000 Blackwell | 0.1158 ms | 0.1371 ms |
+| RTX A6000 | builder smoke 0.2166 ms | 0.3230 ms |
+
+The final `kernel-builder` gate passed on a Vast RTX A6000 with
+`BUILDER_VARIANT=torch210-cxx11-cu128-x86_64-linux`: local pytest `6 passed`,
+decode smoke `0.2166 ms/iter`, builder pytest `4 passed, 2 skipped`, exit
+status `0`.

benchmarks/benchmark_hydra_decode.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+import time
+from pathlib import Path
+
+import torch
+from kernels import get_local_kernel
+
+
+def load_local_kernel(repo: Path):
+    for variant in (repo / "build", repo):
+        if (variant / "metadata.json").exists():
+            return get_local_kernel(variant)
+
+    sys.path.insert(0, str(repo / "torch-ext"))
+    return importlib.import_module("hydra")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Hydra decode microbenchmark")
+    parser.add_argument("--repo", default=".", help="Path to the hydra kernel directory")
+    parser.add_argument("--batch", type=int, default=1)
+    parser.add_argument("--heads", type=int, default=32)
+    parser.add_argument("--kv-heads", type=int, default=8)
+    parser.add_argument("--tokens", type=int, default=8192)
+    parser.add_argument("--head-dim", type=int, default=128)
+    parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--warmup", type=int, default=10)
+    parser.add_argument("--window", type=int, default=0)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if not torch.cuda.is_available():
+        raise SystemExit("CUDA is required")
+
+    kernel = load_local_kernel(Path(args.repo))
+    q = torch.randn(
+        args.batch,
+        args.heads,
+        1,
+        args.head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    k = torch.randn(
+        args.batch,
+        args.kv_heads,
+        args.tokens,
+        args.head_dim,
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    v = torch.randn_like(k)
+
+    window = None if args.window <= 0 else args.window
+    for _ in range(args.warmup):
+        kernel.hydra(q, k, v, sliding_window=window)
+    torch.cuda.synchronize()
+
+    start = time.perf_counter()
+    for _ in range(args.iters):
+        kernel.hydra(q, k, v, sliding_window=window)
+    torch.cuda.synchronize()
+    elapsed = time.perf_counter() - start
+
+    ms = elapsed * 1000.0 / args.iters
+    print(
+        "hydra_decode "
+        f"B={args.batch} H={args.heads} Hkv={args.kv_heads} "
+        f"Tkv={args.tokens} D={args.head_dim} window={window or 0} "
+        f"iters={args.iters} ms_per_iter={ms:.4f}"
+    )
+
+
+if __name__ == "__main__":
+    main()

build.toml

@@ -0,0 +1,7 @@
+[general]
+name = "hydra"
+license = "MIT"
+backends = ["cuda"]
+
+[general.hub]
+repo-id = "Frosty40/hydra"

flake.lock

@@ -0,0 +1,95 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1765121682,
+        "narHash": "sha256-4VBOP18BFeiPkyhy9o4ssBNQEvfvv1kXkasAYd0+rrA=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "65f23138d8d09a92e30f1e5c87611b23ef451bf3",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs"
+      },
+      "locked": {
+        "lastModified": 1775482375,
+        "narHash": "sha256-RUbxfJGs96jwnwSci3+8h08GB2/katOI67yZTCaV+aE=",
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "rev": "dffbce5a048648febb96bbed79f3811ab6d7a577",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernel-builder",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1766341660,
+        "narHash": "sha256-4yG6vx7Dddk9/zh45Y2KM82OaRD4jO3HA9r98ORzysA=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "26861f5606e3e4d1400771b513cc63e5f70151a6",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "ref": "nixos-unstable-small",
+        "repo": "nixpkgs",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,17 @@
+{
+  description = "Flake for the Hydra Hugging Face kernel";
+
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernel-builder";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genFlakeOutputs {
+      path = ./.;
+      rev = self.shortRev or self.dirtyShortRev or (self.lastModifiedDate or "unknown");
+    };
+}

readme_example.py

@@ -0,0 +1,47 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#   "torch",
+#   "triton",
+#   "kernels",
+# ]
+# ///
+
+import os
+from pathlib import Path
+import sys
+
+import torch
+from kernels import get_kernel, get_local_kernel
+
+
+def load_hydra_kernel():
+    if os.environ.get("HYDRA_USE_HUB") == "1":
+        return get_kernel("Frosty40/hydra")
+
+    root = Path(__file__).resolve().parent
+    for variant in (root / "build", root):
+        if (variant / "metadata.json").exists():
+            return get_local_kernel(variant)
+
+    sys.path.insert(0, str(root / "torch-ext"))
+    import hydra
+
+    return hydra
+
+
+def main() -> None:
+    if not torch.cuda.is_available():
+        raise SystemExit("Hydra requires CUDA for this example")
+
+    kernel = load_hydra_kernel()
+    q = torch.randn(1, 32, 1, 128, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 8, 8192, 128, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(1, 8, 8192, 128, device="cuda", dtype=torch.bfloat16)
+
+    out = kernel.hydra(q, k, v)
+    print(f"Hydra decode: {tuple(q.shape)} x {tuple(k.shape)} -> {tuple(out.shape)}")
+
+
+if __name__ == "__main__":
+    main()

tests/conftest.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parents[1]
+TORCH_EXT = ROOT / "torch-ext"
+sys.path.insert(0, str(TORCH_EXT))

tests/test_csr.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("torch")
+pytest.importorskip("triton")
+
+
+def test_dense_causal_csr_cpu_contract():
+    from hydra.csr import build_dense_causal_csr
+
+    row_ptr, col_idx, seq_lens = build_dense_causal_csr(
+        batch_size=1,
+        num_heads=2,
+        seq_len=128,
+        block_size=32,
+        device="cpu",
+    )
+
+    assert row_ptr.shape == (1, 2, 5)
+    assert seq_lens.tolist() == [128]
+    assert row_ptr[0, 0].tolist() == [0, 1, 3, 6, 10]
+    assert col_idx[0, 0].tolist() == [0, 0, 1, 0, 1, 2, 0, 1, 2, 3]
+
+
+def test_sliding_window_csr_keeps_diagonal_last():
+    from hydra.csr import build_sliding_window_csr
+
+    row_ptr, col_idx, _ = build_sliding_window_csr(
+        window=64,
+        seq_len=128,
+        block_size=32,
+        batch_size=1,
+        num_heads=1,
+        device="cpu",
+    )
+
+    rp = row_ptr[0, 0].tolist()
+    ci = col_idx[0, 0].tolist()
+    for q_block in range(4):
+        lo, hi = rp[q_block], rp[q_block + 1]
+        assert ci[hi - 1] == q_block
+        assert all(k < q_block for k in ci[lo : hi - 1])

tests/test_decode_parity.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+pytest.importorskip("triton")
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_decode_matches_sdpa_full_kv():
+    import torch.nn.functional as F
+    import hydra
+
+    torch.manual_seed(0)
+    q = torch.randn(1, 32, 1, 128, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 8, 256, 128, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(1, 8, 256, 128, device="cuda", dtype=torch.bfloat16)
+
+    out = hydra.hydra(q, k, v)
+    k_rep = k.repeat_interleave(4, dim=1)
+    v_rep = v.repeat_interleave(4, dim=1)
+    ref = F.scaled_dot_product_attention(q, k_rep, v_rep, is_causal=False)
+
+    torch.testing.assert_close(out, ref, atol=3e-2, rtol=3e-2)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def test_decode_matches_sdpa_sliding_window():
+    import torch.nn.functional as F
+    import hydra
+
+    torch.manual_seed(1)
+    q = torch.randn(1, 32, 1, 128, device="cuda", dtype=torch.bfloat16)
+    k = torch.randn(1, 8, 257, 128, device="cuda", dtype=torch.bfloat16)
+    v = torch.randn(1, 8, 257, 128, device="cuda", dtype=torch.bfloat16)
+
+    window = 96
+    out = hydra.hydra(q, k, v, sliding_window=window)
+    k_rep = k[:, :, -window:, :].repeat_interleave(4, dim=1)
+    v_rep = v[:, :, -window:, :].repeat_interleave(4, dim=1)
+    ref = F.scaled_dot_product_attention(q, k_rep, v_rep, is_causal=False)
+
+    torch.testing.assert_close(out, ref, atol=3e-2, rtol=3e-2)

tests/test_imports.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("torch")
+pytest.importorskip("triton")
+
+
+def test_public_imports():
+    import hydra
+
+    assert callable(hydra.hydra)
+    assert hydra.hydra_attention is hydra.hydra
+    assert hydra.flash_attn_blackwell is hydra.hydra
+
+
+def test_policy_defaults():
+    from hydra.policy import RuntimePolicy
+
+    policy = RuntimePolicy()
+    assert policy.mode == "off"
+    assert policy.enabled is False

torch-ext/hydra/__init__.py

@@ -0,0 +1,32 @@
+"""Hydra bounded-residency attention package.
+
+This package currently exposes the extracted Triton attention kernels and
+runtime policy layer from the research repo. The public API is deliberately
+small while the benchmark evidence is being consolidated.
+"""
+
+from .api import flash_attn_blackwell, hydra, hydra_attention
+from .csr import build_dense_causal_csr, build_sliding_window_csr
+from .kernel_decode import launch_attn_fwd_decode
+from .policy import (
+    RuntimePolicy,
+    last_policy_decision,
+    policy_history,
+    set_runtime_policy,
+)
+
+__version__ = "0.1.0"
+
+
+__all__ = [
+    "hydra",
+    "hydra_attention",
+    "flash_attn_blackwell",
+    "build_dense_causal_csr",
+    "build_sliding_window_csr",
+    "launch_attn_fwd_decode",
+    "RuntimePolicy",
+    "set_runtime_policy",
+    "last_policy_decision",
+    "policy_history",
+]

torch-ext/hydra/api.py

@@ -0,0 +1,160 @@
+"""Public Python API (perf-tuned variant).
+
+Same public signature as ``hydra.api.hydra``.
+
+Internally:
+
+- Caches the (row_ptr, col_idx, seq_lens) tuple keyed on
+  ``(B, Hq, T, BLOCK_SIZE, window_arg, device, dtype_marker)``. The
+  upstream api re-builds the CSR on every call, which (a) pays a Python
+  cost per call and (b) defeats any data_ptr()-based downstream cache.
+
+Cache invalidation rule
+-----------------------
+A new entry is created whenever any of the following change:
+  - batch size B
+  - num query heads Hq (CSR is broadcast across H so this is part of key)
+  - sequence length T (kernel tile constraint: T % BLOCK_SIZE == 0)
+  - BLOCK_SIZE (module-level constant; included for safety)
+  - window_arg (the effective in-kernel window; 0 for dense)
+  - device (per-CUDA-device pattern; CPU vs CUDA included)
+
+The cache is bounded (default 32 entries). Reaching the cap evicts the
+LRU entry. ``seq_lens`` is included in the cached tuple — by design
+``seq_lens`` is a single ``torch.full`` produced by the CSR builders
+and depends only on B / T, which are part of the key.
+"""
+from __future__ import annotations
+
+import os
+from collections import OrderedDict
+
+import torch
+
+from .csr import build_dense_causal_csr, build_sliding_window_csr
+from .kernel_fwd import BLOCK_SIZE, HEAD_DIM
+from .function import FlashAttnHydraFunction, FlashAttnHydraDecodeFunction
+from .policy import apply_runtime_policy
+
+
+_CSR_CACHE_MAX = int(os.environ.get("HYDRA_CSR_CACHE_MAX", "32"))
+_CSR_CACHE: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor, torch.Tensor]]" = OrderedDict()
+
+
+def _get_csr(
+    B: int,
+    Hq: int,
+    T: int,
+    window_arg: int,
+    device: torch.device,
+    sliding_window: int | None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    key = (
+        B,
+        Hq,
+        T,
+        BLOCK_SIZE,
+        window_arg,
+        sliding_window,
+        device.type,
+        getattr(device, "index", None),
+    )
+    hit = _CSR_CACHE.get(key)
+    if hit is not None:
+        _CSR_CACHE.move_to_end(key)
+        return hit
+
+    if window_arg == 0:
+        row_ptr, col_idx, seq_lens = build_dense_causal_csr(B, Hq, T, BLOCK_SIZE, device)
+    else:
+        row_ptr, col_idx, seq_lens = build_sliding_window_csr(
+            window_arg, T, BLOCK_SIZE, B, Hq, device
+        )
+    _CSR_CACHE[key] = (row_ptr, col_idx, seq_lens)
+    while len(_CSR_CACHE) > _CSR_CACHE_MAX:
+        _CSR_CACHE.popitem(last=False)
+    return row_ptr, col_idx, seq_lens
+
+
+def hydra(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    *,
+    is_causal: bool = True,
+    sliding_window: int | None = None,
+    policy_layer_idx: int | None = None,
+    precision: str = "high",
+) -> torch.Tensor:
+    """Blackwell-tuned causal FlashAttention with GQA and optional sliding window.
+
+    Same semantics as ``hydra.hydra``.
+    """
+    if precision not in {"high", "fast"}:
+        raise ValueError(f"precision must be 'high' or 'fast', got {precision!r}")
+    if precision != "high":
+        raise NotImplementedError("precision='fast' is not wired in this extracted Hydra API yet")
+    if not is_causal:
+        raise NotImplementedError("non-causal attention is not supported in this version")
+    if q.dtype != torch.bfloat16 or k.dtype != torch.bfloat16 or v.dtype != torch.bfloat16:
+        raise ValueError(f"q/k/v must be bf16; got {q.dtype}, {k.dtype}, {v.dtype}")
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(f"q/k/v must be 4D (B, H, T, D); got {q.shape}, {k.shape}, {v.shape}")
+
+    B, Hq, T, D = q.shape
+    if D != HEAD_DIM:
+        raise ValueError(f"head_dim={D} must equal HEAD_DIM={HEAD_DIM}")
+    Hkv = k.shape[1]
+    if Hq % Hkv != 0:
+        raise ValueError(f"H_q={Hq} must be a multiple of H_kv={Hkv} (GQA constraint)")
+
+    decision = apply_runtime_policy(
+        q, k, v,
+        sliding_window=sliding_window,
+        block_size=BLOCK_SIZE,
+        head_dim=HEAD_DIM,
+        layer_idx=policy_layer_idx,
+    )
+    sliding_window = decision.effective_window
+
+    # Decode-step specialization: T_q == 1 with arbitrary T_kv. The decode
+    # kernel does not require T_kv % BLOCK_SIZE == 0 and has no Q-blocking,
+    # so it services HF generation's per-token call without the eager fallback.
+    if T == 1:
+        if sliding_window is not None and sliding_window <= 0:
+            raise ValueError(f"sliding_window must be positive, got {sliding_window}")
+        window_arg = 0 if sliding_window is None else sliding_window
+        return FlashAttnHydraDecodeFunction.apply(q, k, v, window_arg)
+
+    if T % BLOCK_SIZE != 0:
+        raise ValueError(f"T={T} must be a multiple of BLOCK_SIZE={BLOCK_SIZE}")
+
+    if sliding_window is None:
+        window_arg = 0
+    else:
+        if sliding_window <= 0:
+            raise ValueError(f"sliding_window must be positive, got {sliding_window}")
+        # When the window covers the full sequence, this degrades to dense causal;
+        # skip the kernel-side masking work in that case.
+        if sliding_window >= T:
+            window_arg = 0
+        else:
+            window_arg = sliding_window
+
+    row_ptr, col_idx, seq_lens = _get_csr(B, Hq, T, window_arg, q.device, sliding_window)
+
+    return FlashAttnHydraFunction.apply(q, k, v, row_ptr, col_idx, seq_lens, window_arg)
+
+
+def csr_cache_clear() -> None:
+    """Drop all cached (row_ptr, col_idx, seq_lens) entries."""
+    _CSR_CACHE.clear()
+
+
+def csr_cache_info() -> dict:
+    return {"size": len(_CSR_CACHE), "max": _CSR_CACHE_MAX}
+
+
+# Compatibility aliases while the extraction is still being consolidated.
+hydra_attention = hydra
+flash_attn_blackwell = hydra

torch-ext/hydra/csr.py

@@ -0,0 +1,131 @@
+"""CSR (Compressed Sparse Row) builders for the Blackwell FA kernel.
+
+Convention
+----------
+The kernel iterates each Q-block's CSR row as ``ColIdx[ci_lo .. ci_hi)`` where
+the **last entry** is the diagonal block (``col_idx[ci_hi - 1] == q_block_id``)
+and the preceding entries are off-diagonal lower-triangular K-blocks (all
+strictly less than ``q_block_id``, by the causal constraint).
+
+The kernel applies causal masking + hi/lo precision split *within* the
+diagonal block and a standard online-softmax merge for the off-diagonal
+blocks. Builders MUST place the diagonal as the last entry per Q row;
+otherwise the kernel reads the wrong block as "the one needing causal
+masking" and produces silently-wrong attention outputs.
+
+All builders return ``(row_ptr, col_idx, seq_lens)``:
+    row_ptr:  (B, H, num_q_blocks + 1) int32 — CSR row pointers per (B, H).
+    col_idx:  (B, H, total_nnz) int32     — CSR column indices.
+    seq_lens: (B,)            int32       — per-batch sequence length.
+
+The (B, H) broadcast is materialized as a contiguous expand so the kernel can
+index without per-head pointer arithmetic.
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+
+
+def _broadcast_csr(
+    row_ptr_row: list[int],
+    col_idx_row: list[int],
+    batch_size: int,
+    num_heads: int,
+    seq_len: int,
+    device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Lift per-row Python lists to per-(B,H) contiguous int32 tensors."""
+    row_ptr = (
+        torch.tensor(row_ptr_row, device=device, dtype=torch.int32)
+        .view(1, 1, -1)
+        .expand(batch_size, num_heads, -1)
+        .contiguous()
+    )
+    col_idx = (
+        torch.tensor(col_idx_row, device=device, dtype=torch.int32)
+        .view(1, 1, -1)
+        .expand(batch_size, num_heads, -1)
+        .contiguous()
+    )
+    seq_lens = torch.full((batch_size,), seq_len, device=device, dtype=torch.int32)
+    return row_ptr, col_idx, seq_lens
+
+
+def build_dense_causal_csr(
+    batch_size: int,
+    num_heads: int,
+    seq_len: int,
+    block_size: int,
+    device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Dense lower-triangular causal pattern.
+
+    For Q-block i: off-diagonal entries enumerate K-blocks 0..i-1, then the
+    diagonal placeholder i. nnz scales as O(num_q_blocks^2 / 2).
+    """
+    num_q_blocks = math.ceil(seq_len / block_size)
+    row_ptr_row = [0]
+    col_idx_row: list[int] = []
+    for q_block in range(num_q_blocks):
+        col_idx_row.extend(range(q_block))
+        col_idx_row.append(q_block)
+        row_ptr_row.append(len(col_idx_row))
+    return _broadcast_csr(row_ptr_row, col_idx_row, batch_size, num_heads, seq_len, device)
+
+
+def build_sliding_window_csr(
+    window: int,
+    seq_len: int,
+    block_size: int,
+    batch_size: int,
+    num_heads: int,
+    device: torch.device | str,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Causal sliding-window pattern, widened for per-Q-token masking.
+
+    ``window`` is measured in tokens (transformers / Mistral convention):
+    Q-token at position p attends K-tokens [max(0, p - window + 1), p].
+
+    The CSR is built using the **"some Q-token attends"** criterion — a
+    K-block is included if ANY Q-token in the Q-block could attend ANY
+    K-token in it. The kernel applies a per-Q-token in-window mask inside
+    the off-diagonal-leftmost and diagonal blocks, so over-included
+    K-tokens at the leftmost edge get masked to -inf and contribute zero.
+
+    Concretely, for Q-block i, the leftmost-attended K-token across all
+    Q-tokens in the block is ``max(0, i*BS - window + 1)``. The K-block
+    containing that token is its position // BS.
+
+    Edge cases:
+        - window >= seq_len: degrades to dense_causal (full lower triangle).
+        - window == 1: each Q-token only attends itself; diagonal-only with
+          within-block causal mask.
+
+    nnz scales as O(num_q_blocks * ceil(window / BS)) — linear in T for
+    fixed window, which is the point of sliding window.
+    """
+    if window <= 0:
+        raise ValueError(f"window must be positive, got {window}")
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    num_q_blocks = math.ceil(seq_len / block_size)
+    row_ptr_row = [0]
+    col_idx_row: list[int] = []
+    for q_block in range(num_q_blocks):
+        # Smallest K-token position attended by any Q-token in this Q-block.
+        # The leftmost Q-token (position q_block * BS) has the most reach to
+        # the left; its window starts at q_block*BS - window + 1.
+        left_bound = q_block * block_size - window + 1
+        if left_bound <= 0:
+            k_min_block = 0
+        else:
+            # Smallest n with (n+1)*BS - 1 >= left_bound, i.e., n = left_bound // BS.
+            k_min_block = left_bound // block_size
+        # Off-diagonal entries: k_min_block .. q_block - 1
+        col_idx_row.extend(range(k_min_block, q_block))
+        # Diagonal placeholder
+        col_idx_row.append(q_block)
+        row_ptr_row.append(len(col_idx_row))
+    return _broadcast_csr(row_ptr_row, col_idx_row, batch_size, num_heads, seq_len, device)

torch-ext/hydra/function.py

@@ -0,0 +1,75 @@
+"""Autograd ``Function`` wrapper for the Blackwell FA kernel (perf-tuned variant).
+
+Changes vs. upstream ``hydra.function``:
+
+- The transposed-CSR is fetched through ``build_csrT_cached`` (LRU keyed
+  on pattern *content*, not ``data_ptr()``). The upstream cache misses
+  every call today because ``api.hydra`` re-allocates
+  ``row_ptr`` / ``col_idx`` per invocation; the new key fixes that.
+
+- ``rp_T`` / ``ci_T`` are saved through ``ctx.save_for_backward`` so the
+  backward never re-builds them. This matches the upstream behaviour but
+  is now meaningfully cheap on cache hit (~O(num_q_blocks) device-side
+  for the key check).
+"""
+from __future__ import annotations
+
+import torch
+
+from .kernel_fwd import BLOCK_SIZE, launch_attn_fwd
+from .kernel_bwd import build_csrT_cached, launch_attn_bwd
+from .kernel_decode import launch_attn_fwd_decode
+
+
+class FlashAttnHydraFunction(torch.autograd.Function):
+    """Autograd-aware wrapper. See ``api.hydra`` for the public surface."""
+
+    @staticmethod
+    def forward(ctx, q, k, v, row_ptr, col_idx, seq_lens, window):
+        o, lse = launch_attn_fwd(q, k, v, row_ptr, col_idx, seq_lens, window=window)
+        num_q_blocks = q.shape[2] // BLOCK_SIZE
+        rp_T, ci_T = build_csrT_cached(row_ptr, col_idx, num_q_blocks)
+        ctx.save_for_backward(q, k, v, o, lse, row_ptr, col_idx, seq_lens, rp_T, ci_T)
+        ctx.window = window
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        q, k, v, o, lse, row_ptr, col_idx, seq_lens, rp_T, ci_T = ctx.saved_tensors
+        # Both o and do must be contiguous for the delta kernel. The saved o
+        # comes from launch_attn_fwd (torch.empty_like(q) — contiguous if q is)
+        # but transformers can save a view of o into ctx, and `do` arriving
+        # from upstream is often a transpose view from the attention output
+        # reshape. Make both contiguous defensively.
+        dq, dk, dv = launch_attn_bwd(
+            q, k, v, o.contiguous(), do.contiguous(), lse,
+            row_ptr, col_idx, seq_lens,
+            row_ptr_T=rp_T, col_idx_T=ci_T,
+            window=ctx.window,
+        )
+        return dq, dk, dv, None, None, None, None
+
+
+class FlashAttnHydraDecodeFunction(torch.autograd.Function):
+    """Forward-only autograd wrapper for the T_q==1 decode kernel.
+
+    Generation runs under torch.no_grad so backward is intentionally
+    unimplemented; calling .backward() raises with a clear message.
+    """
+
+    @staticmethod
+    def forward(ctx, q, k, v, window):
+        o, lse = launch_attn_fwd_decode(q, k, v, window=window)
+        ctx.save_for_backward(q, k, v, o, lse)
+        ctx.window = window
+        ctx.set_materialize_grads(False)
+        return o
+
+    @staticmethod
+    def backward(ctx, do):
+        raise NotImplementedError(
+            "FlashAttnHydraDecodeFunction has no backward. "
+            "Decode-step (T_q == 1) is forward-only — wrap your call in "
+            "torch.no_grad() (HF generation does this automatically). "
+            "For training, use the prefill kernel with T_q == T_kv."
+        )

torch-ext/hydra/kernel_bwd.py

@@ -0,0 +1,743 @@
+"""Backward Triton kernels for Blackwell-tuned causal FlashAttention.
+
+Perf-tuned variant of ``hydra.kernel_bwd`` that
+
+  - keeps the Triton kernels byte-for-byte identical to the upstream
+    versions (so the gradient-correctness story is unchanged),
+  - replaces ``build_csrT``'s CPU + numpy round-trip with a pure-torch /
+    GPU-resident equivalent (``build_csrT_gpu``) that produces the same
+    output and stays on the input device,
+  - exposes a module-level LRU cache (``build_csrT_cached``) keyed on
+    pattern shape rather than ``data_ptr()`` so per-call rebuilds of
+    ``row_ptr`` / ``col_idx`` (as the api does today) still hit the
+    cache,
+  - clears the host-side ``delta = (do.float() * o.float()).sum(-1)``
+    cast spew: the existing launcher computes ``delta`` with two fp32
+    upcasts and a stride-collapsing ``.contiguous()`` every call; we keep
+    the math but drop the redundant ``.contiguous()`` (the result of
+    ``.sum(-1)`` on a contiguous 4D tensor is already contiguous).
+
+Public API is preserved: ``launch_attn_bwd`` has the same signature, and
+``build_csrT`` is kept as a thin wrapper that dispatches to the pure-GPU
+builder when possible (and falls back to the original CPU+numpy path
+when the input is on CPU, so the CPU equivalence test exercises the
+same code path).
+"""
+from __future__ import annotations
+
+import math
+import os
+from collections import OrderedDict
+
+import torch
+import triton
+import triton.language as tl
+
+from .kernel_fwd import BLOCK_SIZE, HEAD_DIM
+from .kernel_delta import launch_compute_delta
+
+
+_BWD_NUM_WARPS = int(os.environ.get("HYDRA_BWD_NUM_WARPS", "4"))
+_BWD_NUM_STAGES = int(os.environ.get("HYDRA_BWD_NUM_STAGES", "1"))
+_DISABLE_AUTOTUNE = int(os.environ.get("HYDRA_DISABLE_AUTOTUNE", "0"))
+
+_CSRT_CACHE_MAX = int(os.environ.get("HYDRA_CSRT_CACHE_MAX", "32"))
+
+
+def _autotune_configs() -> list[triton.Config]:
+    configs: list[triton.Config] = []
+    for num_warps in (2, 4, 8):
+        for num_stages in (1, 2):
+            configs.append(triton.Config({}, num_warps=num_warps, num_stages=num_stages))
+    return configs
+
+
+def _kernel_decorator(jit_kernel):
+    if _DISABLE_AUTOTUNE:
+        return jit_kernel
+    return triton.autotune(
+        configs=_autotune_configs(),
+        # Include WINDOW: sliding-window patterns have 2-4× fewer K-blocks per
+        # Q-block than dense, so optimal warps/stages plausibly differ. Keying
+        # them separately costs one extra autotune sweep per (T, window) pair.
+        key=["T_MAX", "D", "NUM_HEADS", "NUM_KV_HEADS", "WINDOW"],
+    )(jit_kernel)
+
+
+# ----------------------------- csrT builders -----------------------------
+
+
+def _build_csrT_cpu_reference(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int):
+    """Original CPU + numpy implementation, kept for reference / fallback.
+
+    Identical semantics to the upstream ``build_csrT`` in
+    ``hydra.kernel_bwd``: per (B, H), enumerate which
+    Q-blocks reference each K-block as an OFF-diagonal attendee. The
+    diagonal placeholder at ``col_idx[ci_hi - 1]`` is excluded.
+    """
+    import numpy as np
+    rp = row_ptr.detach().cpu().numpy()
+    ci = col_idx.detach().cpu().numpy()
+    B, H = rp.shape[0], rp.shape[1]
+    rp_flat = rp.reshape(B * H, -1)
+    ci_flat = ci.reshape(B * H, -1)
+    rpt_rows: list[list[int]] = []
+    cit_rows: list[list[int]] = []
+    for bh in range(B * H):
+        bucket: list[list[int]] = [[] for _ in range(num_blocks)]
+        rp_row = rp_flat[bh]
+        ci_row = ci_flat[bh]
+        for i in range(num_blocks):
+            lo = int(rp_row[i])
+            hi = int(rp_row[i + 1])
+            for p in range(lo, max(lo, hi - 1)):
+                bucket[int(ci_row[p])].append(i)
+        offs = [0]
+        vals: list[int] = []
+        for k in range(num_blocks):
+            vals.extend(bucket[k])
+            offs.append(len(vals))
+        rpt_rows.append(offs)
+        cit_rows.append(vals)
+    max_nnz = max(1, max(len(v) for v in cit_rows))
+    cit_pad = np.zeros((B * H, max_nnz), dtype=np.int32)
+    for bh, v in enumerate(cit_rows):
+        if v:
+            cit_pad[bh, : len(v)] = v
+    rpt_np = np.asarray(rpt_rows, dtype=np.int32)
+    rp_T = torch.from_numpy(rpt_np).to(row_ptr.device).reshape(B, H, num_blocks + 1).contiguous()
+    ci_T = torch.from_numpy(cit_pad).to(row_ptr.device).reshape(B, H, max_nnz).contiguous()
+    return rp_T, ci_T
+
+
+def _is_broadcast_pattern(row_ptr: torch.Tensor, col_idx: torch.Tensor) -> bool:
+    """Cheap check: does the (B, H) pattern collapse to a single shared row?
+
+    The CSR builders in ``hydra.csr`` always produce a
+    broadcast-then-contiguous pattern. We exploit that to compute the
+    transposed CSR once and ``.expand().contiguous()`` it back to
+    (B, H). For safety we verify by comparing row_ptr/col_idx against
+    the (0, 0) head.
+    """
+    B, H = row_ptr.shape[0], row_ptr.shape[1]
+    if B == 1 and H == 1:
+        return True
+    rp0 = row_ptr[0, 0]
+    ci0 = col_idx[0, 0]
+    # equal_to expects same-shape comparand; use a single torch.equal per
+    # axis. This costs O(B*H*nnz) but is GPU-resident and ~1us for typical
+    # shapes — far cheaper than the .cpu()/.numpy() round-trip.
+    return bool(torch.all(row_ptr == rp0).item()) and bool(torch.all(col_idx == ci0).item())
+
+
+def _build_csrT_single_row(rp_row: torch.Tensor, ci_row: torch.Tensor, num_blocks: int):
+    """Pure-torch transposed-CSR for a single (B=H=1) row.
+
+    Produces ``(rpt_row, cit_row)`` of dtype int32 on the same device as
+    the inputs. The ordering of q_block_ids within each K-row matches
+    the CPU reference: ascending q_block_id.
+    """
+    device = rp_row.device
+    rp_row = rp_row.to(torch.int64)
+    ci_row = ci_row.to(torch.int64)
+
+    # Off-diagonal "off" slices per Q-row: [lo, hi - 1). Total off-diag
+    # entries summed across Q-blocks equals nnz - num_blocks (one diag
+    # per Q-block). Build a flat (q_block_id, k_block_id) edge list.
+    counts = (rp_row[1:] - rp_row[:-1] - 1).clamp(min=0)  # off-diag count per Q-block
+    num_offdiag = int(counts.sum().item())
+    if num_offdiag == 0:
+        rp_T_row = torch.zeros(num_blocks + 1, dtype=torch.int32, device=device)
+        ci_T_row = torch.zeros(1, dtype=torch.int32, device=device)
+        return rp_T_row, ci_T_row, 1
+
+    # q_idx[e] = which Q-block this edge belongs to.
+    q_idx = torch.repeat_interleave(
+        torch.arange(num_blocks, dtype=torch.int64, device=device),
+        counts,
+    )
+    # k_idx[e] = which K-block this edge points at. We need the slice
+    # ci_row[lo : hi - 1] for each Q-row, concatenated. Build a flat
+    # index into ci_row by starting at lo for each Q-block and adding the
+    # within-row offset (0, 1, 2, ...).
+    starts = rp_row[:-1]  # lo per Q-block
+    edge_within = torch.arange(num_offdiag, dtype=torch.int64, device=device) - torch.repeat_interleave(
+        torch.cat([
+            torch.zeros(1, dtype=torch.int64, device=device),
+            counts.cumsum(0)[:-1],
+        ]),
+        counts,
+    )
+    edge_pos = torch.repeat_interleave(starts, counts) + edge_within
+    k_idx = ci_row[edge_pos]
+
+    # Sort edges by k_idx (stable so q_idx within a bucket stays ascending).
+    sort_idx = torch.argsort(k_idx, stable=True)
+    k_idx_sorted = k_idx[sort_idx]
+    q_idx_sorted = q_idx[sort_idx]
+
+    # rp_T[k+1] = count of edges with k_idx <= k. Use bincount over [0, num_blocks).
+    per_k_count = torch.bincount(k_idx_sorted, minlength=num_blocks)
+    rp_T_row = torch.zeros(num_blocks + 1, dtype=torch.int64, device=device)
+    rp_T_row[1:] = per_k_count.cumsum(0)
+
+    max_nnz = max(1, num_offdiag)
+    ci_T_row = torch.zeros(max_nnz, dtype=torch.int32, device=device)
+    ci_T_row[:num_offdiag] = q_idx_sorted.to(torch.int32)
+    return rp_T_row.to(torch.int32), ci_T_row, max_nnz
+
+
+def build_csrT_gpu(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int):
+    """Pure-device transposed-CSR builder.
+
+    Equivalent to ``build_csrT`` but never touches CPU. When the input
+    pattern is the broadcast-across-(B,H) form (the common case), only
+    one head's worth of work is done and the result is expanded.
+
+    Falls back to a per-(B,H) torch implementation for the non-broadcast
+    case (no current api path produces that, but keep it safe).
+    """
+    if row_ptr.dim() != 3 or col_idx.dim() != 3:
+        raise ValueError(f"expected 3D row_ptr/col_idx; got {row_ptr.shape}, {col_idx.shape}")
+    B, H = row_ptr.shape[0], row_ptr.shape[1]
+
+    if _is_broadcast_pattern(row_ptr, col_idx):
+        rp_T_row, ci_T_row, max_nnz = _build_csrT_single_row(row_ptr[0, 0], col_idx[0, 0], num_blocks)
+        rp_T = rp_T_row.view(1, 1, num_blocks + 1).expand(B, H, num_blocks + 1).contiguous()
+        ci_T = ci_T_row.view(1, 1, max_nnz).expand(B, H, max_nnz).contiguous()
+        return rp_T, ci_T
+
+    # Non-broadcast path: per-(B,H) computation. We loop in Python over
+    # the B*H heads but each head's work stays on-device. For B*H up to
+    # a few thousand this is still much cheaper than .cpu().numpy().
+    out_rp_rows: list[torch.Tensor] = []
+    out_ci_rows: list[torch.Tensor] = []
+    max_nnz_seen = 1
+    for bh in range(B * H):
+        bi, hi = bh // H, bh % H
+        rp_T_row, ci_T_row, mn = _build_csrT_single_row(row_ptr[bi, hi], col_idx[bi, hi], num_blocks)
+        out_rp_rows.append(rp_T_row)
+        out_ci_rows.append(ci_T_row)
+        max_nnz_seen = max(max_nnz_seen, mn)
+
+    rp_T = torch.stack(out_rp_rows, dim=0).view(B, H, num_blocks + 1).contiguous()
+    # Right-pad to max_nnz_seen.
+    ci_T = torch.zeros(B * H, max_nnz_seen, dtype=torch.int32, device=row_ptr.device)
+    for bh, row in enumerate(out_ci_rows):
+        ci_T[bh, : row.shape[0]] = row
+    ci_T = ci_T.view(B, H, max_nnz_seen).contiguous()
+    return rp_T, ci_T
+
+
+def build_csrT(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int):
+    """Drop-in replacement for the upstream ``build_csrT``.
+
+    Routes to the pure-GPU builder on CUDA inputs and to the
+    CPU+numpy reference on CPU inputs (so the equivalence test exercises
+    the same numerical path the kernel sees).
+    """
+    if row_ptr.is_cuda:
+        return build_csrT_gpu(row_ptr, col_idx, num_blocks)
+    return _build_csrT_cpu_reference(row_ptr, col_idx, num_blocks)
+
+
+# ----------------------------- csrT cache --------------------------------
+
+
+# Two-level cache.
+#
+#   - _CSRT_CACHE_FAST is keyed on (row_ptr.data_ptr(), col_idx.data_ptr(),
+#     shapes, device, num_blocks). When the public api caches the CSR tensors,
+#     the data_ptrs are stable across calls and this is an O(1) Python dict probe — no
+#     CUDA sync, no D2H copy.
+#
+#   - _CSRT_CACHE_CONTENT is keyed on (num_blocks, shape, first-row
+#     content). Used as a fallback when the fast key misses (e.g. a user
+#     who builds CSRs fresh each call). The first-row content read costs
+#     one (num_blocks+1) + first-row-nnz int32 D2H sync — still vastly
+#     cheaper than the full CPU+numpy round-trip on the build side.
+#
+# On a content-cache hit we also promote into the fast cache so the next
+# call with the same data_ptr is free.
+_CSRT_CACHE_FAST: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict()
+_CSRT_CACHE_CONTENT: "OrderedDict[tuple, tuple[torch.Tensor, torch.Tensor]]" = OrderedDict()
+
+
+def _fast_key(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int) -> tuple:
+    return (
+        row_ptr.data_ptr(),
+        col_idx.data_ptr(),
+        num_blocks,
+        tuple(row_ptr.shape),
+        tuple(col_idx.shape),
+        row_ptr.device.type,
+        getattr(row_ptr.device, "index", None),
+    )
+
+
+def _content_key(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int) -> tuple:
+    """Content-addressed key. Reads the (0, 0) head's row to cover the
+    broadcast-CSR case (the only one the api emits today).
+
+    Invalidation rule
+    -----------------
+    A new entry is created whenever any of the following change:
+      - num_blocks
+      - row_ptr.shape / col_idx.shape
+      - device type / index
+      - the (0, 0) head's row_ptr / col_idx values
+
+    Non-broadcast inputs that share the (0, 0) row but differ on other
+    (B, H) entries would alias under this key. The api never produces
+    such inputs, but ``launch_attn_bwd`` users with hand-built CSRs
+    should call ``csrT_cache_clear()`` between distinct non-broadcast
+    patterns.
+    """
+    rp0 = tuple(row_ptr[0, 0].to(torch.int64).tolist())
+    ci0 = tuple(col_idx[0, 0].to(torch.int64).tolist())
+    return (
+        num_blocks,
+        tuple(row_ptr.shape),
+        tuple(col_idx.shape),
+        row_ptr.device.type,
+        getattr(row_ptr.device, "index", None),
+        rp0,
+        ci0,
+    )
+
+
+def _bound(cache: "OrderedDict") -> None:
+    while len(cache) > _CSRT_CACHE_MAX:
+        cache.popitem(last=False)
+
+
+def build_csrT_cached(row_ptr: torch.Tensor, col_idx: torch.Tensor, num_blocks: int):
+    """Return a (rp_T, ci_T) pair for the given CSR, hitting the LRU cache when possible.
+
+    Thread-safety: this is not thread-safe across CUDA streams. Callers
+    that share the same CSR across streams should manually clone the
+    returned tensors. In practice the autograd backward is called on the
+    same stream as the forward that produced the CSR, so this is fine.
+    """
+    fk = _fast_key(row_ptr, col_idx, num_blocks)
+    hit = _CSRT_CACHE_FAST.get(fk)
+    if hit is not None:
+        _CSRT_CACHE_FAST.move_to_end(fk)
+        return hit
+    ck = _content_key(row_ptr, col_idx, num_blocks)
+    hit = _CSRT_CACHE_CONTENT.get(ck)
+    if hit is not None:
+        _CSRT_CACHE_CONTENT.move_to_end(ck)
+        _CSRT_CACHE_FAST[fk] = hit
+        _bound(_CSRT_CACHE_FAST)
+        return hit
+    rp_T, ci_T = build_csrT(row_ptr, col_idx, num_blocks)
+    _CSRT_CACHE_FAST[fk] = (rp_T, ci_T)
+    _CSRT_CACHE_CONTENT[ck] = (rp_T, ci_T)
+    _bound(_CSRT_CACHE_FAST)
+    _bound(_CSRT_CACHE_CONTENT)
+    return rp_T, ci_T
+
+
+def csrT_cache_clear() -> None:
+    """Drop all cached csrT entries. Useful for test isolation."""
+    _CSRT_CACHE_FAST.clear()
+    _CSRT_CACHE_CONTENT.clear()
+
+
+def csrT_cache_info() -> dict:
+    """Return basic cache stats. ``size`` is the number of *distinct*
+    csrT patterns currently cached (the content-keyed dict)."""
+    return {
+        "size": len(_CSRT_CACHE_CONTENT),
+        "fast_size": len(_CSRT_CACHE_FAST),
+        "max": _CSRT_CACHE_MAX,
+    }
+
+
+# ----------------------------- Triton kernels ----------------------------
+# These are kept local so the extracted package is self-contained.
+
+
+@triton.jit
+def _hydra_bwd_dq_jit(
+    Q, K, V, LSE,
+    dO,
+    Di_in,
+    dQ,
+    RowPtr, ColIdx, SeqLens,
+    stride_cih,
+    T_MAX: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    SCALE: tl.constexpr,
+    BS: tl.constexpr,
+    D: tl.constexpr,
+    WINDOW: tl.constexpr,
+    CONTIGUOUS_OFFDIAG: tl.constexpr = 0,
+):
+    stride_h: tl.constexpr = T_MAX * D
+    stride_t: tl.constexpr = D
+    stride_rph: tl.constexpr = (T_MAX // BS) + 1
+    stride_lh: tl.constexpr = T_MAX
+    NUM_Q_BLOCKS: tl.constexpr = T_MAX // BS
+
+    pid = tl.program_id(0)
+    bh_id = pid // NUM_Q_BLOCKS
+    q_block_id = pid % NUM_Q_BLOCKS
+
+    rep: tl.constexpr = NUM_HEADS // NUM_KV_HEADS
+    b_id = bh_id // NUM_HEADS
+    hq_id = bh_id % NUM_HEADS
+    hkv_id = hq_id // rep
+    kv_bh_id = b_id * NUM_KV_HEADS + hkv_id
+
+    seq_len = tl.load(SeqLens + b_id)
+    q_start = q_block_id * BS
+    if q_start >= seq_len:
+        return
+
+    offs_tok = q_start + tl.arange(0, BS)
+    offs_d = tl.arange(0, D)
+    q_mask = offs_tok < seq_len
+
+    Q_ptr = Q + bh_id * stride_h
+    q_bf16 = tl.load(Q_ptr + offs_tok[:, None] * stride_t + offs_d[None, :])
+
+    dO_ptr = dO + bh_id * stride_h
+    dO_i = tl.load(dO_ptr + offs_tok[:, None] * stride_t + offs_d[None, :])
+
+    LSE_ptr = LSE + bh_id * stride_lh
+    LSE_i = tl.load(LSE_ptr + offs_tok)
+
+    rp_base = RowPtr + bh_id * stride_rph + q_block_id
+    ci_lo = tl.load(rp_base)
+    ci_hi = tl.load(rp_base + 1)
+
+    K_ptr = K + kv_bh_id * stride_h
+    V_ptr = V + kv_bh_id * stride_h
+    CI_ptr = ColIdx + bh_id * stride_cih
+
+    offs_k_tile = tl.arange(0, BS)
+    boundary = (q_start + BS) > seq_len
+
+    Di = tl.load(Di_in + bh_id * stride_lh + offs_tok, mask=q_mask, other=0.0)
+
+    dQ_i = tl.zeros([BS, D], dtype=tl.float32)
+
+    if ci_hi > ci_lo:
+        k_start_d = q_block_id * BS
+        offs_k_d = k_start_d + offs_k_tile
+        if boundary:
+            k_mask_d = offs_k_d < seq_len
+            k_bf16_d = tl.load(K_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :], mask=k_mask_d[:, None], other=0.0)
+            v_bf16_d = tl.load(V_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :], mask=k_mask_d[:, None], other=0.0)
+            s_d = tl.dot(q_bf16, tl.trans(k_bf16_d), out_dtype=tl.float32) * SCALE
+            causal = offs_tok[:, None] >= offs_k_d[None, :]
+            allowed = causal & k_mask_d[None, :]
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k_d[None, :]) < WINDOW
+                allowed = allowed & in_window
+            s_d = tl.where(allowed, s_d, float("-inf"))
+        else:
+            k_bf16_d = tl.load(K_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+            v_bf16_d = tl.load(V_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+            s_d = tl.dot(q_bf16, tl.trans(k_bf16_d), out_dtype=tl.float32) * SCALE
+            causal = offs_tok[:, None] >= offs_k_d[None, :]
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k_d[None, :]) < WINDOW
+                s_d = tl.where(causal & in_window, s_d, float("-inf"))
+            else:
+                s_d = tl.where(causal, s_d, float("-inf"))
+
+        P_d = tl.exp(s_d - LSE_i[:, None])
+        dP_d = tl.dot(dO_i, tl.trans(v_bf16_d), out_dtype=tl.float32)
+        dS_d = (P_d * (dP_d - Di[:, None])) * SCALE
+        dQ_i += tl.dot(dS_d.to(tl.bfloat16), k_bf16_d, out_dtype=tl.float32)
+
+    # See kernel_fwd: when CONTIGUOUS_OFFDIAG, hoist the CSR load.
+    if CONTIGUOUS_OFFDIAG:
+        k_block_id_start = tl.load(CI_ptr + ci_lo)
+        for ci in range(ci_lo, ci_hi - 1):
+            k_block_id = k_block_id_start + (ci - ci_lo)
+            k_start = k_block_id * BS
+            offs_k = k_start + offs_k_tile
+            k_bf16 = tl.load(K_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            v_bf16 = tl.load(V_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            s = tl.dot(q_bf16, tl.trans(k_bf16), out_dtype=tl.float32) * SCALE
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k[None, :]) < WINDOW
+                s = tl.where(in_window, s, float("-inf"))
+            P = tl.exp(s - LSE_i[:, None])
+            dP = tl.dot(dO_i, tl.trans(v_bf16), out_dtype=tl.float32)
+            dS = (P * (dP - Di[:, None])) * SCALE
+            dQ_i += tl.dot(dS.to(tl.bfloat16), k_bf16, out_dtype=tl.float32)
+    else:
+        for ci in range(ci_lo, ci_hi - 1):
+            k_block_id = tl.load(CI_ptr + ci)
+            k_start = k_block_id * BS
+            offs_k = k_start + offs_k_tile
+            k_bf16 = tl.load(K_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            v_bf16 = tl.load(V_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            s = tl.dot(q_bf16, tl.trans(k_bf16), out_dtype=tl.float32) * SCALE
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k[None, :]) < WINDOW
+                s = tl.where(in_window, s, float("-inf"))
+            P = tl.exp(s - LSE_i[:, None])
+            dP = tl.dot(dO_i, tl.trans(v_bf16), out_dtype=tl.float32)
+            dS = (P * (dP - Di[:, None])) * SCALE
+            dQ_i += tl.dot(dS.to(tl.bfloat16), k_bf16, out_dtype=tl.float32)
+
+    dQ_ptr = dQ + bh_id * stride_h
+    tl.store(dQ_ptr + offs_tok[:, None] * stride_t + offs_d[None, :], dQ_i.to(tl.bfloat16), mask=q_mask[:, None])
+
+
+@triton.jit
+def _hydra_bwd_dkv_jit(
+    Q, K, V, LSE,
+    dO,
+    Di_in,
+    dK, dV,
+    RowPtrT, ColIdxT, SeqLens,
+    stride_cith,
+    T_MAX: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    SCALE: tl.constexpr,
+    BS: tl.constexpr,
+    D: tl.constexpr,
+    WINDOW: tl.constexpr,
+    CONTIGUOUS_OFFDIAG: tl.constexpr = 0,
+):
+    stride_h: tl.constexpr = T_MAX * D
+    stride_t: tl.constexpr = D
+    stride_rpth: tl.constexpr = (T_MAX // BS) + 1
+    stride_lh: tl.constexpr = T_MAX
+    NUM_K_BLOCKS: tl.constexpr = T_MAX // BS
+
+    pid = tl.program_id(0)
+    kv_bh_id = pid // NUM_K_BLOCKS
+    k_block_id = pid % NUM_K_BLOCKS
+
+    rep: tl.constexpr = NUM_HEADS // NUM_KV_HEADS
+    b_id = kv_bh_id // NUM_KV_HEADS
+    hkv_id = kv_bh_id % NUM_KV_HEADS
+    q_head_start = hkv_id * rep
+
+    seq_len = tl.load(SeqLens + b_id)
+    k_start = k_block_id * BS
+    if k_start >= seq_len:
+        return
+
+    offs_d = tl.arange(0, D)
+    offs_tok_k = k_start + tl.arange(0, BS)
+    k_mask = offs_tok_k < seq_len
+
+    K_ptr = K + kv_bh_id * stride_h
+    V_ptr = V + kv_bh_id * stride_h
+
+    k_bf16 = tl.load(K_ptr + offs_tok_k[:, None] * stride_t + offs_d[None, :], mask=k_mask[:, None], other=0.0)
+    v_bf16 = tl.load(V_ptr + offs_tok_k[:, None] * stride_t + offs_d[None, :], mask=k_mask[:, None], other=0.0)
+
+    dK_acc = tl.zeros([BS, D], dtype=tl.float32)
+    dV_acc = tl.zeros([BS, D], dtype=tl.float32)
+
+    for rep_idx in range(rep):
+        q_head_id = q_head_start + rep_idx
+        q_bh_id = b_id * NUM_HEADS + q_head_id
+
+        Q_ptr = Q + q_bh_id * stride_h
+        dO_ptr = dO + q_bh_id * stride_h
+        LSE_ptr = LSE + q_bh_id * stride_lh
+        Di_ptr = Di_in + q_bh_id * stride_lh
+
+        q_start_diag = k_block_id * BS
+        offs_tok_q = q_start_diag + tl.arange(0, BS)
+        q_mask_diag = offs_tok_q < seq_len
+        q_bf16_d = tl.load(Q_ptr + offs_tok_q[:, None] * stride_t + offs_d[None, :], mask=q_mask_diag[:, None], other=0.0)
+        dO_d = tl.load(dO_ptr + offs_tok_q[:, None] * stride_t + offs_d[None, :], mask=q_mask_diag[:, None], other=0.0)
+        LSE_d = tl.load(LSE_ptr + offs_tok_q, mask=q_mask_diag, other=0.0)
+        Di_d = tl.load(Di_ptr + offs_tok_q, mask=q_mask_diag, other=0.0)
+
+        s_d = tl.dot(q_bf16_d, tl.trans(k_bf16), out_dtype=tl.float32) * SCALE
+        causal_d = offs_tok_q[:, None] >= offs_tok_k[None, :]
+        allowed_d = causal_d & k_mask[None, :] & q_mask_diag[:, None]
+        if WINDOW > 0:
+            in_window_d = (offs_tok_q[:, None] - offs_tok_k[None, :]) < WINDOW
+            allowed_d = allowed_d & in_window_d
+        s_d = tl.where(allowed_d, s_d, float("-inf"))
+        P_d = tl.exp(s_d - LSE_d[:, None])
+        dP_d = tl.dot(dO_d, tl.trans(v_bf16), out_dtype=tl.float32)
+        dS_d = (P_d * (dP_d - Di_d[:, None])) * SCALE
+
+        dK_acc += tl.dot(tl.trans(dS_d.to(tl.bfloat16)), q_bf16_d, out_dtype=tl.float32)
+        dV_acc += tl.dot(tl.trans(P_d.to(tl.bfloat16)), dO_d, out_dtype=tl.float32)
+
+        rpt_base = RowPtrT + q_bh_id * stride_rpth + k_block_id
+        cit_lo = tl.load(rpt_base)
+        cit_hi = tl.load(rpt_base + 1)
+        CIT_ptr = ColIdxT + q_bh_id * stride_cith
+
+        # See kernel_fwd: when CONTIGUOUS_OFFDIAG, hoist the transposed-CSR
+        # load. For dense, K-block k attends Q-blocks k+1..num_q_blocks-1
+        # (contiguous); for sliding-window, k attends a contiguous Q-range
+        # too. The transposed-CSR builder (build_csrT) preserves ascending
+        # q_block_id order within each K-row, so the run is contiguous.
+        if CONTIGUOUS_OFFDIAG:
+            q_block_id_start_o = tl.load(CIT_ptr + cit_lo)
+            for ci in range(cit_lo, cit_hi):
+                q_block_id = q_block_id_start_o + (ci - cit_lo)
+                q_start = q_block_id * BS
+                offs_tok_q2 = q_start + tl.arange(0, BS)
+                q_mask2 = offs_tok_q2 < seq_len
+                q_bf16_o = tl.load(Q_ptr + offs_tok_q2[:, None] * stride_t + offs_d[None, :], mask=q_mask2[:, None], other=0.0)
+                dO_o = tl.load(dO_ptr + offs_tok_q2[:, None] * stride_t + offs_d[None, :], mask=q_mask2[:, None], other=0.0)
+                LSE_o = tl.load(LSE_ptr + offs_tok_q2, mask=q_mask2, other=0.0)
+                Di_o = tl.load(Di_ptr + offs_tok_q2, mask=q_mask2, other=0.0)
+
+                s_o = tl.dot(q_bf16_o, tl.trans(k_bf16), out_dtype=tl.float32) * SCALE
+                allowed_o = q_mask2[:, None] & k_mask[None, :]
+                if WINDOW > 0:
+                    in_window_o = (offs_tok_q2[:, None] - offs_tok_k[None, :]) < WINDOW
+                    allowed_o = allowed_o & in_window_o
+                s_o = tl.where(allowed_o, s_o, float("-inf"))
+                P_o = tl.exp(s_o - LSE_o[:, None])
+                dP_o = tl.dot(dO_o, tl.trans(v_bf16), out_dtype=tl.float32)
+                dS_o = (P_o * (dP_o - Di_o[:, None])) * SCALE
+
+                dK_acc += tl.dot(tl.trans(dS_o.to(tl.bfloat16)), q_bf16_o, out_dtype=tl.float32)
+                dV_acc += tl.dot(tl.trans(P_o.to(tl.bfloat16)), dO_o, out_dtype=tl.float32)
+        else:
+            for ci in range(cit_lo, cit_hi):
+                q_block_id = tl.load(CIT_ptr + ci)
+                q_start = q_block_id * BS
+                offs_tok_q2 = q_start + tl.arange(0, BS)
+                q_mask2 = offs_tok_q2 < seq_len
+                q_bf16_o = tl.load(Q_ptr + offs_tok_q2[:, None] * stride_t + offs_d[None, :], mask=q_mask2[:, None], other=0.0)
+                dO_o = tl.load(dO_ptr + offs_tok_q2[:, None] * stride_t + offs_d[None, :], mask=q_mask2[:, None], other=0.0)
+                LSE_o = tl.load(LSE_ptr + offs_tok_q2, mask=q_mask2, other=0.0)
+                Di_o = tl.load(Di_ptr + offs_tok_q2, mask=q_mask2, other=0.0)
+
+                s_o = tl.dot(q_bf16_o, tl.trans(k_bf16), out_dtype=tl.float32) * SCALE
+                allowed_o = q_mask2[:, None] & k_mask[None, :]
+                if WINDOW > 0:
+                    in_window_o = (offs_tok_q2[:, None] - offs_tok_k[None, :]) < WINDOW
+                    allowed_o = allowed_o & in_window_o
+                s_o = tl.where(allowed_o, s_o, float("-inf"))
+                P_o = tl.exp(s_o - LSE_o[:, None])
+                dP_o = tl.dot(dO_o, tl.trans(v_bf16), out_dtype=tl.float32)
+                dS_o = (P_o * (dP_o - Di_o[:, None])) * SCALE
+
+                dK_acc += tl.dot(tl.trans(dS_o.to(tl.bfloat16)), q_bf16_o, out_dtype=tl.float32)
+                dV_acc += tl.dot(tl.trans(P_o.to(tl.bfloat16)), dO_o, out_dtype=tl.float32)
+
+    dK_ptr = dK + kv_bh_id * stride_h + offs_tok_k[:, None] * stride_t + offs_d[None, :]
+    dV_ptr = dV + kv_bh_id * stride_h + offs_tok_k[:, None] * stride_t + offs_d[None, :]
+    tl.store(dK_ptr, dK_acc.to(tl.bfloat16), mask=k_mask[:, None])
+    tl.store(dV_ptr, dV_acc.to(tl.bfloat16), mask=k_mask[:, None])
+
+
+_hydra_bwd_dq_kernel = _kernel_decorator(_hydra_bwd_dq_jit)
+_hydra_bwd_dkv_kernel = _kernel_decorator(_hydra_bwd_dkv_jit)
+
+
+# ----------------------------- launcher ----------------------------------
+
+
+def launch_attn_bwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    o: torch.Tensor,
+    do: torch.Tensor,
+    lse: torch.Tensor,
+    row_ptr: torch.Tensor,
+    col_idx: torch.Tensor,
+    seq_lens: torch.Tensor,
+    row_ptr_T: torch.Tensor | None = None,
+    col_idx_T: torch.Tensor | None = None,
+    window: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Launch dQ then dK/dV kernels.
+
+    Identical signature / semantics to the upstream ``launch_attn_bwd``.
+    If ``row_ptr_T`` / ``col_idx_T`` are not provided, builds them via
+    the cached pure-GPU builder so repeated calls with the same pattern
+    pay a single one-time build cost.
+    """
+    batch_size, num_heads, t_max, head_dim = q.shape
+    num_kv_heads = k.shape[1]
+    if t_max % BLOCK_SIZE != 0:
+        raise ValueError(f"T ({t_max}) must be a multiple of BLOCK_SIZE ({BLOCK_SIZE})")
+    if head_dim != HEAD_DIM:
+        raise ValueError(f"head_dim ({head_dim}) must equal HEAD_DIM ({HEAD_DIM})")
+    if num_heads % num_kv_heads != 0:
+        raise ValueError(f"num_heads ({num_heads}) must be divisible by num_kv_heads ({num_kv_heads})")
+
+    batch_heads_q = batch_size * num_heads
+    batch_heads_kv = batch_size * num_kv_heads
+    num_q_blocks = t_max // BLOCK_SIZE
+
+    dq = torch.empty_like(q)
+    dk = torch.empty_like(k)
+    dv = torch.empty_like(v)
+
+    # delta[b, h, t] = sum_d O[b, h, t, d] * dO[b, h, t, d]
+    # Single-pass Triton kernel avoids the ~3 fp32 transients (do.float(),
+    # o.float(), their product) the host pipeline allocates — ~256 MB
+    # steady-state HBM saved at Qwen3-8B T=8192. Bounded reorder error
+    # vs torch's .sum(-1) is ~1.5e-5, three orders below bf16 quantum.
+    delta = launch_compute_delta(o, do)
+    delta_2d = delta.view(batch_heads_q, t_max)
+    lse_2d = lse.view(batch_heads_q, t_max)
+
+    if row_ptr_T is None or col_idx_T is None:
+        row_ptr_T, col_idx_T = build_csrT_cached(row_ptr, col_idx, num_q_blocks)
+
+    stride_cih = col_idx.shape[2] if col_idx.ndim == 3 else col_idx.shape[1]
+    stride_cith = col_idx_T.shape[2] if col_idx_T.ndim == 3 else col_idx_T.shape[1]
+    scale = 1.0 / math.sqrt(head_dim)
+
+    grid_q = (num_q_blocks * batch_heads_q,)
+    grid_kv = (num_q_blocks * batch_heads_kv,)
+
+    common_kwargs = dict(
+        T_MAX=t_max,
+        NUM_HEADS=num_heads,
+        NUM_KV_HEADS=num_kv_heads,
+        SCALE=scale,
+        BS=BLOCK_SIZE,
+        D=head_dim,
+        WINDOW=int(window),
+        # Both build_dense_causal_csr and build_sliding_window_csr emit
+        # off-diag K-block IDs as a contiguous run; the transposed CSR
+        # builder preserves ascending q_block_id within each K-row, which
+        # is likewise contiguous for these patterns. Lets the kernels
+        # hoist the per-iteration CSR scalar GMEM load.
+        CONTIGUOUS_OFFDIAG=1,
+    )
+    if _DISABLE_AUTOTUNE:
+        common_kwargs["num_warps"] = _BWD_NUM_WARPS
+        common_kwargs["num_stages"] = _BWD_NUM_STAGES
+
+    _hydra_bwd_dq_kernel[grid_q](
+        q, k, v, lse_2d,
+        do,
+        delta_2d,
+        dq,
+        row_ptr, col_idx, seq_lens,
+        stride_cih,
+        **common_kwargs,
+    )
+
+    _hydra_bwd_dkv_kernel[grid_kv](
+        q, k, v, lse_2d,
+        do,
+        delta_2d,
+        dk, dv,
+        row_ptr_T, col_idx_T, seq_lens,
+        stride_cith,
+        **common_kwargs,
+    )
+
+    return dq, dk, dv

torch-ext/hydra/kernel_decode.py

@@ -0,0 +1,298 @@
+"""Forward Triton kernel for the decode step: Q.seq_len == 1.
+
+Specialization
+--------------
+Generation hot path: one new Q token attends to a full KV cache of length
+T_kv. The prefill kernel in ``kernel_fwd.py`` requires T % BLOCK_SIZE == 0
+and pays for Q-blocking + causal CSR iteration that is meaningless when
+there is exactly one query row. This kernel collapses to:
+
+    program grid = (B * H_q,)              # one program per query head
+    Q broadcast-replicated to a [MM, D] tile (MM=16, the tcore M minimum)
+    K/V streamed in [BK, D] tiles across the visible K range
+
+Online softmax with a single Q row reduces to scalar (m, l) and a [1, D]
+accumulator (we lift it to [MM, D] with broadcast so Triton's tl.dot —
+which requires M >= 16 on Blackwell — sees a tensor-core-sized matmul).
+Only the first row of the [MM, D] output is actually distinct; we keep
+the broadcast cost (~16× duplicate work) because the V matmul is the
+heavy cost and dominates regardless. The alternative — manually emitting
+fma loops — would not use tensor cores and would be far slower.
+
+No causal mask is needed because the single Q token is at absolute
+position T_kv - 1 and every visible K token is <= that position.
+
+Sliding window
+--------------
+When ``WINDOW > 0`` and ``WINDOW < T_kv``, K is restricted to the last
+WINDOW positions: ``k_start = T_kv - WINDOW``. The launcher computes
+``k_start`` and the number of K tiles; the kernel iterates that range
+unconditionally and masks the tail tile where ``offs_k >= T_kv``.
+
+GQA
+---
+Each H_q head reads the same K/V slab indexed by
+``hkv = hq // (NUM_HEADS // NUM_KV_HEADS)``.
+
+Precision
+---------
+bf16 inputs, fp32 accumulator, bf16 output. The kernel casts ``p`` to bf16
+once per K-tile before the ``p @ V`` matmul. No hi/lo precision split:
+without a causal mask there is no single "diagonal" block where the cast
+loss is concentrated; for T_kv up to ~32K this stays well inside bf16's
+useful range and the test below confirms a max abs diff < 3e-2 against
+SDPA.
+"""
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+
+# K-tile size for the decode kernel. Independent of the prefill kernel's
+# BLOCK_SIZE because the decode kernel never reads CSR; it just streams.
+BLOCK_K = int(os.environ.get("HYDRA_DECODE_BLOCK_K", "64"))
+HEAD_DIM = int(os.environ.get("HYDRA_HEAD_DIM", "128"))
+
+_DEC_NUM_WARPS = int(os.environ.get("HYDRA_DECODE_NUM_WARPS", "4"))
+_DEC_NUM_STAGES = int(os.environ.get("HYDRA_DECODE_NUM_STAGES", "2"))
+_DISABLE_AUTOTUNE = int(os.environ.get("HYDRA_DISABLE_AUTOTUNE", "0"))
+
+
+def _autotune_configs() -> list[triton.Config]:
+    """Decode-friendly configs. Q is a single row so warp pressure is low;
+    favour stages=2-3 to overlap K/V loads with the small matmuls.
+    """
+    configs: list[triton.Config] = []
+    for num_warps in (2, 4, 8):
+        for num_stages in (2, 3, 4):
+            configs.append(triton.Config({}, num_warps=num_warps, num_stages=num_stages))
+    return configs
+
+
+def _fixed_config() -> triton.Config:
+    return triton.Config({}, num_warps=_DEC_NUM_WARPS, num_stages=_DEC_NUM_STAGES)
+
+
+def _kernel_decorator(jit_kernel):
+    if _DISABLE_AUTOTUNE:
+        return jit_kernel
+    return triton.autotune(
+        configs=_autotune_configs(),
+        # T_KV is a runtime arg (not constexpr) and is intentionally NOT in the
+        # key — every generation token changes T_kv by 1, and including it
+        # forced per-token recompiles (~5-10s each, 200x slowdown observed).
+        key=["D", "NUM_HEADS", "NUM_KV_HEADS"],
+    )(jit_kernel)
+
+
+@triton.jit
+def _hydra_decode_jit(
+    Q,                # (B, H_q, 1, D) bf16
+    K, V,             # (B, H_kv, T_KV, D) bf16
+    O,                # (B, H_q, 1, D) bf16
+    LSE,              # (B, H_q, 1) fp32  (kept for parity / future use)
+    T_KV,                    # runtime int — used only as the mask bound; NOT constexpr to avoid per-T_kv recompiles during generation
+    stride_kvbh_kv,          # runtime stride for (B*Hkv) axis of K/V in elements (lets StaticCache pass non-T_KV*D strides)
+    stride_t_kv,             # runtime stride for the T axis of K/V in elements
+    NUM_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    SCALE: tl.constexpr,
+    BK: tl.constexpr,
+    D: tl.constexpr,
+    MM: tl.constexpr,        # Q-replication dim (16 — tcore M-minimum).
+    K_START: tl.constexpr,   # leftmost token attended (>=0). Sliding-window left edge.
+    NUM_K_TILES: tl.constexpr,
+):
+    # Q strides: per-(B, H_q) slab has 1*D = D elements (Q is always contig (B,Hq,1,D)).
+    # K/V strides: passed in as runtime args so StaticCache (B, Hkv, T_max, D) works
+    # — kernel previously hardcoded stride_kvbh = T_KV * D which is wrong for non-T_KV strides.
+    stride_qbh: tl.constexpr = D
+    stride_kvbh = stride_kvbh_kv
+    stride_t = stride_t_kv
+
+    pid = tl.program_id(0)
+    rep: tl.constexpr = NUM_HEADS // NUM_KV_HEADS
+    b_id = pid // NUM_HEADS
+    hq_id = pid % NUM_HEADS
+    hkv_id = hq_id // rep
+
+    q_bh = b_id * NUM_HEADS + hq_id
+    kv_bh = b_id * NUM_KV_HEADS + hkv_id
+
+    offs_d = tl.arange(0, D)
+    offs_mm = tl.arange(0, MM)
+    offs_k_in_tile = tl.arange(0, BK)
+
+    # Fold SCALE * log2(e) into Q at load so the per-K-tile score does not
+    # pay one mul per element.
+    LOG2E: tl.constexpr = 1.4426950408889634
+    SCALE_2: tl.constexpr = SCALE * LOG2E
+    Q_ptr = Q + q_bh * stride_qbh
+
+    # Tensor cores need M >= 16 in tl.dot. Q has only one valid row, so we
+    # replicate it MM times along the M axis. All MM result rows of any
+    # downstream matmul are then identical, and we only keep row 0.
+    q_row = tl.load(Q_ptr + offs_d).to(tl.float32) * SCALE_2        # [D] fp32
+    q_bf16 = tl.broadcast_to(q_row[None, :], [MM, D]).to(tl.bfloat16)  # [MM, D]
+
+    m_i = tl.full([MM], float("-inf"), dtype=tl.float32)
+    l_i = tl.zeros([MM], dtype=tl.float32)
+    acc = tl.zeros([MM, D], dtype=tl.float32)
+
+    K_ptr = K + kv_bh * stride_kvbh
+    V_ptr = V + kv_bh * stride_kvbh
+
+    # Iterate K tiles over the visible window [K_START, T_KV).
+    # Tile t covers tokens [K_START + t*BK, K_START + (t+1)*BK).
+    for t in range(0, NUM_K_TILES):
+        k_base = K_START + t * BK
+        offs_k = k_base + offs_k_in_tile                 # [BK]
+        k_mask = offs_k < T_KV                           # [BK]
+
+        # Masked load: OOB columns get 0; we still mask the score to -inf
+        # so its softmax weight is 0.
+        k_tile = tl.load(
+            K_ptr + offs_k[:, None] * stride_t + offs_d[None, :],
+            mask=k_mask[:, None], other=0.0,
+        )                                                # [BK, D] bf16
+        v_tile = tl.load(
+            V_ptr + offs_k[:, None] * stride_t + offs_d[None, :],
+            mask=k_mask[:, None], other=0.0,
+        )                                                # [BK, D] bf16
+
+        # Score: [MM, D] @ [D, BK] -> [MM, BK] in fp32. All MM rows identical.
+        s = tl.dot(q_bf16, tl.trans(k_tile), out_dtype=tl.float32)    # [MM, BK]
+        s = tl.where(k_mask[None, :], s, float("-inf"))
+
+        # Online softmax merge. Identical per row, so the [MM]-vector state
+        # stays uniform; we keep the broadcast for matmul-shape consistency.
+        s_max = tl.max(s, axis=1)                        # [MM]
+        m_new = tl.maximum(m_i, s_max)
+        alpha = tl.exp2(m_i - m_new)                     # [MM]
+        p = tl.exp2(s - m_new[:, None])                  # [MM, BK]
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+        acc = tl.dot(p.to(tl.bfloat16), v_tile, acc=acc * alpha[:, None], out_dtype=tl.float32)
+        m_i = m_new
+
+    LN2: tl.constexpr = 0.6931471805599453
+    # All MM rows of acc / l_i / m_i are identical (they each saw the same
+    # broadcast Q row through the entire loop). Reduce to the row-0 values
+    # for the final store. ``tl.sum(..., axis=0) / MM`` extracts the common
+    # value as a side effect of averaging identical replicas — this is
+    # numerically equivalent to picking row 0 but stays inside Triton's
+    # supported reduction ops without needing a slice.
+    l_safe = tl.where(l_i > 0, l_i, 1.0)                 # [MM]
+    o_tile = acc / l_safe[:, None]                       # [MM, D] fp32, all rows equal
+    INV_MM: tl.constexpr = 1.0 / MM
+    # Average the MM identical replicas to extract the row-0 result.
+    # tl.sum(axis=0) over [MM, D] -> [D]; over [MM] -> scalar wrapped as [1]
+    # after a no-op broadcast through tl.arange masking.
+    o_row = tl.sum(o_tile, axis=0) * INV_MM              # [D] fp32
+    o_row_bf16 = o_row.to(tl.bfloat16)                   # [D] bf16
+    lse_full = tl.where(l_i > 0, (m_i + tl.log2(l_safe)) * LN2, float("-inf"))  # [MM]
+    # Build a [1]-shaped LSE: sum the [MM] vector then divide; the result is a
+    # Triton scalar — wrap via tl.full to get an explicit [1]-shape value.
+    lse_scalar_val = tl.sum(lse_full, axis=0) * INV_MM   # scalar fp32
+    lse_out_vec = tl.full([1], 0.0, dtype=tl.float32) + lse_scalar_val  # [1]
+
+    # Flat-store the single output row [D] at O_ptr + offs_d. No mask needed:
+    # offs_d covers exactly the D-element output buffer for this (B, Hq).
+    O_ptr = O + q_bh * stride_qbh
+    tl.store(O_ptr + offs_d, o_row_bf16)
+
+    # LSE: store the [1]-shaped value at offset q_bh.
+    LSE_ptr = LSE + q_bh
+    tl.store(LSE_ptr + tl.arange(0, 1), lse_out_vec)
+
+
+_hydra_decode_kernel = _kernel_decorator(_hydra_decode_jit)
+
+
+def launch_attn_fwd_decode(
+    q: torch.Tensor,        # (B, H_q, 1, D) bf16
+    k: torch.Tensor,        # (B, H_kv, T_kv, D) bf16
+    v: torch.Tensor,        # (B, H_kv, T_kv, D) bf16
+    window: int = 0,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Launch the decode-step forward kernel.
+
+    Args:
+        q: (B, H_q, 1, D) bf16. T_q must be exactly 1.
+        k: (B, H_kv, T_kv, D) bf16.
+        v: (B, H_kv, T_kv, D) bf16.
+        window: 0 disables sliding window (attend all K). When > 0, only the
+            last ``window`` K tokens are attended: K range = [T_kv - window,
+            T_kv). The query is at position T_kv - 1, so this matches the
+            standard sliding-window-causal semantics used at prefill.
+
+    Returns:
+        o:   (B, H_q, 1, D) bf16
+        lse: (B, H_q, 1)    fp32
+
+    Notes:
+        Assumes K/V are contiguous in (B, H_kv, T_kv, D). T_kv does NOT need
+        to be a multiple of BLOCK_K; the last K-tile is masked.
+    """
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(f"q/k/v must be 4D (B,H,T,D); got {q.shape} {k.shape} {v.shape}")
+    B, Hq, Tq, D = q.shape
+    if Tq != 1:
+        raise ValueError(f"decode kernel requires Tq == 1; got Tq={Tq}")
+    if k.shape[0] != B or v.shape[0] != B:
+        raise ValueError(f"batch mismatch q={q.shape} k={k.shape} v={v.shape}")
+    if k.shape != v.shape:
+        raise ValueError(f"k and v must have identical shapes; got {k.shape} vs {v.shape}")
+    if k.shape[3] != D:
+        raise ValueError(f"head_dim mismatch q={q.shape[-1]} k={k.shape[-1]}")
+    Hkv = k.shape[1]
+    Tkv = k.shape[2]
+    if Hq % Hkv != 0:
+        raise ValueError(f"H_q ({Hq}) must be a multiple of H_kv ({Hkv})")
+    if D != HEAD_DIM:
+        raise ValueError(f"head_dim ({D}) must equal HEAD_DIM ({HEAD_DIM})")
+    if Tkv <= 0:
+        raise ValueError(f"T_kv must be positive; got {Tkv}")
+
+    # Sliding window: clamp to [0, Tkv]. window <= 0 or window >= Tkv -> full.
+    if window is None or window <= 0 or window >= Tkv:
+        k_start = 0
+    else:
+        k_start = Tkv - window
+    num_k_visible = Tkv - k_start
+    num_k_tiles = (num_k_visible + BLOCK_K - 1) // BLOCK_K
+
+    o = torch.empty_like(q)
+    lse = torch.empty((B, Hq, 1), device=q.device, dtype=torch.float32)
+
+    grid = (B * Hq,)
+    # Pull actual strides from K (V must match shape). Lets StaticCache pass
+    # (B, Hkv, T_max, D) buffers — previously hardcoded stride_kvbh = Tkv*D
+    # would silently read garbage when T_max != Tkv.
+    common_kwargs = dict(
+        T_KV=Tkv,
+        stride_kvbh_kv=k.stride(1),
+        stride_t_kv=k.stride(2),
+        NUM_HEADS=Hq,
+        NUM_KV_HEADS=Hkv,
+        SCALE=1.0 / math.sqrt(D),
+        BK=BLOCK_K,
+        D=D,
+        MM=16,     # tensor-core minimum M for tl.dot on Blackwell bf16.
+        K_START=int(k_start),
+        NUM_K_TILES=int(num_k_tiles),
+    )
+    if _DISABLE_AUTOTUNE:
+        common_kwargs["num_warps"] = _DEC_NUM_WARPS
+        common_kwargs["num_stages"] = _DEC_NUM_STAGES
+
+    _hydra_decode_kernel[grid](
+        q, k, v,
+        o, lse,
+        **common_kwargs,
+    )
+    return o, lse

torch-ext/hydra/kernel_delta.py

@@ -0,0 +1,188 @@
+"""Single-pass Triton delta kernel.
+
+Replaces the host-side ``(do.float() * o.float()).sum(-1)`` dance in
+``hydra.kernel_bwd.launch_attn_bwd`` with a fused
+GPU-resident kernel that reads ``o`` and ``do`` once (bf16, the same
+loads the existing autograd chain already pays for) and writes a
+single ``(B, Hq, T)`` fp32 ``delta`` tensor.
+
+What this saves vs the host-side computation (Qwen3-8B-ish prefill,
+B=1, H_q=32, T=8192, D=128):
+
+  - ``do.float()`` 128 MB fp32 transient: gone
+  - ``o.float()``  128 MB fp32 transient: gone
+  - ``do.float() * o.float()`` 128 MB fp32 product: gone
+  - net delta tensor still allocated: 1 MB fp32  (B*Hq*T*4)
+
+For an at-rest measurement the peak transient eliminated is ~384 MB;
+in the steady-state allocator the saving conservatively measured by
+the prior agent's accounting is ~256 MB (after one of the upcasts is
+reaped). Either way the dedicated kernel is the right move — the
+``o``/``do`` bytes are already on-chip in the autograd path, we just
+fold the multiply+reduce into one pass.
+
+Math:
+
+  delta[b, h, t] = sum_{d=0..D-1} o[b, h, t, d] * do[b, h, t, d]
+
+bf16 inputs, fp32 accumulator, fp32 output. The reduction order is
+Triton's parallel ``tl.sum`` tree across the D axis (per tile-row).
+Compared to torch's left-to-right ``.sum(-1)`` the absolute error
+bound is roughly ``D * eps_fp32 * max|o*do|`` which is ~1.5e-5 at
+attention-typical scales — well within bf16's ~4e-3 quantum.
+
+Layout:
+
+  - ``o``, ``do``: ``(B, Hq, T, D)`` bf16, contiguous.
+  - ``delta``:    ``(B, Hq, T)``   fp32, contiguous.
+  - One program per ``(b*Hq + h, q_block_id)`` work item, computing
+    ``BLOCK_SIZE`` rows of delta with full ``D`` in-tile reduction.
+  - Grid: ``(B * Hq * num_q_blocks,)`` flattened.
+
+Arbitrary T: the kernel masks the tail tile (``offs_tok < T``); the
+launcher does not require ``T % BLOCK_SIZE == 0``.
+"""
+from __future__ import annotations
+
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+from .kernel_fwd import BLOCK_SIZE, HEAD_DIM
+
+
+_DELTA_NUM_WARPS = int(os.environ.get("HYDRA_DELTA_NUM_WARPS", "4"))
+_DELTA_NUM_STAGES = int(os.environ.get("HYDRA_DELTA_NUM_STAGES", "1"))
+_DISABLE_AUTOTUNE = int(os.environ.get("HYDRA_DISABLE_AUTOTUNE", "0"))
+
+
+def _autotune_configs() -> list[triton.Config]:
+    """Same warps/stages sweep used by kernel_fwd; D and BS are fixed."""
+    configs: list[triton.Config] = []
+    for num_warps in (1, 2, 4):
+        for num_stages in (1, 2):
+            configs.append(triton.Config({}, num_warps=num_warps, num_stages=num_stages))
+    return configs
+
+
+def _kernel_decorator(jit_kernel):
+    if _DISABLE_AUTOTUNE:
+        return jit_kernel
+    return triton.autotune(
+        configs=_autotune_configs(),
+        key=["T_MAX", "D"],
+    )(jit_kernel)
+
+
+@triton.jit
+def _compute_delta_jit(
+    O,         # (B, Hq, T, D) bf16
+    dO,        # (B, Hq, T, D) bf16
+    Delta,     # (B, Hq, T)    fp32
+    T_MAX: tl.constexpr,
+    BS: tl.constexpr,
+    D: tl.constexpr,
+):
+    """Compute ``Delta[b, h, t] = sum_d O[b, h, t, d] * dO[b, h, t, d]``.
+
+    One program handles a single ``(b*Hq + h, q_block_id)`` tile of
+    ``BS`` rows, reducing the full ``D``-axis in fp32 in-register.
+
+    Grid: ``(batch_heads * num_q_blocks,)``.
+    """
+    # Per-head row strides match the (B, Hq, T, D) / (B, Hq, T) layout
+    # the contiguous-tensor invariant gives us. Encoding them as
+    # constexpr lets the compiler fold the address arithmetic.
+    stride_h: tl.constexpr = T_MAX * D
+    stride_t: tl.constexpr = D
+    stride_lh: tl.constexpr = T_MAX
+    NUM_Q_BLOCKS: tl.constexpr = (T_MAX + BS - 1) // BS
+
+    pid = tl.program_id(0)
+    bh_id = pid // NUM_Q_BLOCKS
+    q_block_id = pid % NUM_Q_BLOCKS
+
+    q_start = q_block_id * BS
+    offs_tok = q_start + tl.arange(0, BS)
+    offs_d = tl.arange(0, D)
+    # Mask trailing tile when T is not a multiple of BS.
+    row_mask = offs_tok < T_MAX
+
+    # Load the (BS, D) o and do tiles in bf16. Out-of-bound rows are
+    # zero-filled so their contribution to the (masked-out) sum is
+    # exactly 0 and doesn't pollute neighbour rows under the tl.sum tree.
+    O_ptr = O + bh_id * stride_h
+    dO_ptr = dO + bh_id * stride_h
+    addr = offs_tok[:, None] * stride_t + offs_d[None, :]
+    o_bf16 = tl.load(O_ptr + addr, mask=row_mask[:, None], other=0.0)
+    do_bf16 = tl.load(dO_ptr + addr, mask=row_mask[:, None], other=0.0)
+
+    # bf16 -> fp32 -> elementwise mul -> reduce along D.
+    # The cast-then-mul order matches the host-side path
+    # ``do.float() * o.float()`` exactly (no intermediate bf16 product).
+    prod = o_bf16.to(tl.float32) * do_bf16.to(tl.float32)
+    di = tl.sum(prod, axis=-1)  # (BS,) fp32
+
+    Delta_ptr = Delta + bh_id * stride_lh
+    tl.store(Delta_ptr + offs_tok, di, mask=row_mask)
+
+
+_compute_delta_kernel = _kernel_decorator(_compute_delta_jit)
+
+
+def launch_compute_delta(o: torch.Tensor, do: torch.Tensor) -> torch.Tensor:
+    """Compute ``delta = sum_d o * do`` as a fused single-pass GPU kernel.
+
+    Parameters
+    ----------
+    o, do : ``(B, Hq, T, D)`` bf16 tensors. Must be contiguous and on the
+            same CUDA device. ``T`` is arbitrary (the tail tile is
+            masked); ``D`` must equal ``HEAD_DIM``.
+
+    Returns
+    -------
+    delta : ``(B, Hq, T)`` fp32 tensor, same device. Allocated fresh.
+
+    Numerical equivalence to ``(do.float() * o.float()).sum(-1)``:
+
+      Both compute ``sum_d (bf16_to_fp32(o[d]) * bf16_to_fp32(do[d]))``
+      with ``D == 128`` summands per (b, h, t) row. Triton's ``tl.sum``
+      uses a parallel reduction tree while torch's ``.sum(-1)`` is a
+      left-to-right scan. The reordering error is bounded by
+      ``D * eps_fp32 * max|o*do|`` ~ ``128 * 1.2e-7 * O(1)`` ~ ``1.5e-5``,
+      well under bf16's ~4e-3 quantum.
+    """
+    if o.shape != do.shape:
+        raise ValueError(f"o and do must have the same shape; got {o.shape} vs {do.shape}")
+    if o.dim() != 4:
+        raise ValueError(f"o must be 4D (B, Hq, T, D); got shape {o.shape}")
+    if o.dtype != torch.bfloat16 or do.dtype != torch.bfloat16:
+        raise ValueError(f"o, do must be bf16; got {o.dtype}, {do.dtype}")
+    if o.device != do.device:
+        raise ValueError(f"o and do must be on the same device; got {o.device} vs {do.device}")
+    if not o.is_contiguous() or not do.is_contiguous():
+        raise ValueError("o and do must be contiguous")
+
+    B, Hq, T, D = o.shape
+    if D != HEAD_DIM:
+        raise ValueError(f"head_dim ({D}) must equal HEAD_DIM ({HEAD_DIM})")
+
+    delta = torch.empty(B, Hq, T, dtype=torch.float32, device=o.device)
+
+    batch_heads = B * Hq
+    num_q_blocks = (T + BLOCK_SIZE - 1) // BLOCK_SIZE
+    grid = (batch_heads * num_q_blocks,)
+
+    kwargs = dict(
+        T_MAX=T,
+        BS=BLOCK_SIZE,
+        D=D,
+    )
+    if _DISABLE_AUTOTUNE:
+        kwargs["num_warps"] = _DELTA_NUM_WARPS
+        kwargs["num_stages"] = _DELTA_NUM_STAGES
+
+    _compute_delta_kernel[grid](o, do, delta, **kwargs)
+    return delta

torch-ext/hydra/kernel_fwd.py

@@ -0,0 +1,379 @@
+"""Cross-arch-tuned forward kernel.
+
+This module contains the active forward Triton kernel used by the extracted
+Hydra package. The JIT kernel body and launcher descend from the iter3
+research kernel, with the cross-architecture tuning hooks kept explicit near
+the top of the file.
+
+Why
+---
+The iter3 autotune sweep is a 9-cell cartesian
+``(num_warps ∈ {2,4,8}) × (num_stages ∈ {1,2,3})`` over a fixed
+``BLOCK_SIZE = 64``. Tuning at import time per
+``torch.cuda.get_device_capability()`` lets us:
+
+1. Lower default ``BLOCK_SIZE`` from 64 → 32 on every measured arch
+   (sm_86/89/120/121). Cross-arch sweep at T=4096 dense fwdbwd found
+   BS=32 winning on every architecture we measured. The smaller K/V
+   tile reduces shared-memory pressure and lets more CTAs run
+   concurrently per SM, which dominates on Ampere/Ada cards
+   (~100 KB smem/SM cap) and is neutral-to-positive on Blackwell.
+
+2. Replace the 9-cell autotune sweep with a 1-3 cell list keyed on
+   compute capability. For arches we've measured we ship the
+   sweep-winning ``(num_warps, num_stages)`` plus one or two
+   safe-fallback configs (so the autotuner can't get stuck if a
+   particular load pattern at runtime regresses on the headline cell).
+   For arches we have NOT measured (anything not in our table) we
+   fall back to the original 9-cell sweep so the autotuner can find
+   a good config on first use.
+
+Both changes preserve the launcher API and the kernel signature.
+No correctness change: ``BS`` is still a constexpr; the kernel just
+gets specialised at a different value.
+
+Per-arch defaults are provisional tuning seeds, not benchmark claims. They
+control autotune order and can be overridden with the environment variables
+below. Submission-facing benchmark claims must come from checked-in artifacts,
+not this header.
+
+Do not cite the defaults below as evidence of speedup or hardware coverage.
+
+Override hooks
+--------------
+Three env vars short-circuit everything:
+  - HYDRA_BLOCK_SIZE - wins over the per-arch default
+  - HYDRA_DISABLE_AUTOTUNE=1 + HYDRA_NUM_WARPS
+    + HYDRA_NUM_STAGES — wins over autotune
+  - HYDRA_AUTOTUNE_MODE=full — forces the 9-cell sweep
+    even on a known arch (useful for re-tuning new wheels of triton).
+"""
+from __future__ import annotations
+
+import math
+import os
+
+import torch
+import triton
+import triton.language as tl
+
+
+# ---------------- Per-arch defaults --------------------------------------
+
+# Per-(major,minor) compute-capability -> (BLOCK_SIZE, [(num_warps, num_stages), ...])
+#
+# The autotune list is ordered with the SWEEP-WINNING config first so that
+# the autotuner picks it on the first launch (which is also the only launch
+# in many production deploys: the cache key includes T_MAX/D so a stable
+# shape produces exactly one autotune cycle).
+#
+# Tail entries are conservative fallbacks: they exist so the autotuner has
+# at least one safe config to fall back to if a specific runtime workload
+# regresses on the headline cell (e.g. a different T that we did not sweep).
+_PER_ARCH: dict[tuple[int, int], tuple[int, list[tuple[int, int]]]] = {
+    # GB10 (Blackwell sm_121).
+    (12, 1): (32, [(2, 3), (4, 2)]),
+    # Pro 6000 (Blackwell sm_120). NOT MEASURED in iter; default to the
+    # sm_121 pick — they share most of the relevant resources. Adjust
+    # after running the sweep on a free Pro 6000.
+    (12, 0): (32, [(2, 3), (4, 2)]),
+    # Ada (sm_89): 4080, 4070 Ti.
+    (8, 9): (32, [(4, 1), (4, 2)]),
+    # Ampere consumer (sm_86): 3090, 3080, 3060.
+    (8, 6): (32, [(2, 1), (4, 2)]),
+}
+
+# The original 9-cell sweep — used when no per-arch entry exists, and also
+# when the user opts in via HYDRA_AUTOTUNE_MODE=full.
+_FULL_SWEEP: list[tuple[int, int]] = [
+    (W, S) for W in (2, 4, 8) for S in (1, 2, 3)
+]
+
+
+def _device_capability_or_none() -> tuple[int, int] | None:
+    """Return ``(major, minor)`` for cuda:0, or ``None`` if no CUDA."""
+    try:
+        if torch.cuda.is_available():
+            return tuple(torch.cuda.get_device_capability(0))  # type: ignore[return-value]
+    except Exception:
+        pass
+    return None
+
+
+def _default_block_size_from_arch() -> int:
+    """Pick BLOCK_SIZE based on the current CUDA device's compute capability.
+
+    Fall back to 64 (iter3's value) if we don't have a measured entry, so
+    no measured arch regresses below its iter3 baseline.
+    """
+    cc = _device_capability_or_none()
+    if cc is not None and cc in _PER_ARCH:
+        return _PER_ARCH[cc][0]
+    return 64
+
+
+def _arch_aware_autotune_configs() -> list[triton.Config]:
+    """Per-arch shrunk autotune list, or the full 9-cell sweep if unknown."""
+    if os.environ.get("HYDRA_AUTOTUNE_MODE", "").lower() == "full":
+        ws_pairs = _FULL_SWEEP
+    else:
+        cc = _device_capability_or_none()
+        if cc is not None and cc in _PER_ARCH:
+            ws_pairs = _PER_ARCH[cc][1]
+        else:
+            ws_pairs = _FULL_SWEEP
+    return [triton.Config({}, num_warps=W, num_stages=S) for (W, S) in ws_pairs]
+
+
+# Env var override wins over per-arch default (preserves the iter3 escape hatch).
+BLOCK_SIZE = int(os.environ.get("HYDRA_BLOCK_SIZE", str(_default_block_size_from_arch())))
+HEAD_DIM = int(os.environ.get("HYDRA_HEAD_DIM", "128"))
+
+_FWD_NUM_WARPS = int(os.environ.get("HYDRA_NUM_WARPS", "4"))
+_FWD_NUM_STAGES = int(os.environ.get("HYDRA_NUM_STAGES", "2"))
+_DISABLE_AUTOTUNE = int(os.environ.get("HYDRA_DISABLE_AUTOTUNE", "0"))
+
+
+def _autotune_configs() -> list[triton.Config]:
+    return _arch_aware_autotune_configs()
+
+
+def _kernel_decorator(jit_kernel):
+    if _DISABLE_AUTOTUNE:
+        return jit_kernel
+    return triton.autotune(
+        configs=_autotune_configs(),
+        key=["T_MAX", "D", "NUM_HEADS", "NUM_KV_HEADS", "ASSUME_FULL"],
+    )(jit_kernel)
+
+
+@triton.jit
+def _hydra_fwd_jit(
+    Q, K, V,
+    O, LSE,
+    RowPtr, ColIdx, SeqLens,
+    stride_cih,
+    T_MAX: tl.constexpr,
+    NUM_HEADS: tl.constexpr,
+    NUM_KV_HEADS: tl.constexpr,
+    SCALE: tl.constexpr,
+    BS: tl.constexpr,
+    D: tl.constexpr,
+    WINDOW: tl.constexpr,
+    CONTIGUOUS_OFFDIAG: tl.constexpr = 0,
+    ASSUME_FULL: tl.constexpr = 0,
+):
+    stride_h: tl.constexpr = T_MAX * D
+    stride_t: tl.constexpr = D
+    stride_lh: tl.constexpr = T_MAX
+    stride_rph: tl.constexpr = (T_MAX // BS) + 1
+    NUM_Q_BLOCKS: tl.constexpr = T_MAX // BS
+
+    pid = tl.program_id(0)
+    bh_id = pid // NUM_Q_BLOCKS
+    q_block_id = pid % NUM_Q_BLOCKS
+
+    rep: tl.constexpr = NUM_HEADS // NUM_KV_HEADS
+    b_id = bh_id // NUM_HEADS
+    hq_id = bh_id % NUM_HEADS
+    hkv_id = hq_id // rep
+    kv_bh_id = b_id * NUM_KV_HEADS + hkv_id
+
+    if ASSUME_FULL:
+        seq_len = T_MAX  # constexpr-friendly: no SeqLens load, no early-return
+    else:
+        seq_len = tl.load(SeqLens + b_id)
+
+    q_start = q_block_id * BS
+
+    offs_tok = q_start + tl.arange(0, BS)
+    offs_d = tl.arange(0, D)
+
+    if ASSUME_FULL:
+        pass  # no early-return path needed
+    else:
+        if q_start >= seq_len:
+            O_ptr_e = O + bh_id * stride_h
+            tl.store(O_ptr_e + offs_tok[:, None] * stride_t + offs_d[None, :],
+                     tl.zeros([BS, D], dtype=tl.bfloat16))
+            LSE_ptr_e = LSE + bh_id * stride_lh
+            tl.store(LSE_ptr_e + offs_tok, tl.full([BS], float("-inf"), dtype=tl.float32))
+            return
+
+    LOG2E: tl.constexpr = 1.4426950408889634
+    SCALE_2: tl.constexpr = SCALE * LOG2E
+    Q_ptr = Q + bh_id * stride_h
+    q_bf16 = (tl.load(Q_ptr + offs_tok[:, None] * stride_t + offs_d[None, :]).to(tl.float32) * SCALE_2).to(tl.bfloat16)
+
+    m_i = tl.full([BS], float("-inf"), dtype=tl.float32)
+    l_i = tl.zeros([BS], dtype=tl.float32)
+    acc = tl.zeros([BS, D], dtype=tl.float32)
+
+    rp_base = RowPtr + bh_id * stride_rph + q_block_id
+    ci_lo = tl.load(rp_base)
+    ci_hi = tl.load(rp_base + 1)
+
+    K_ptr = K + kv_bh_id * stride_h
+    V_ptr = V + kv_bh_id * stride_h
+    CI_ptr = ColIdx + bh_id * stride_cih
+
+    offs_d_arange = tl.arange(0, BS)
+
+    if ci_hi > ci_lo:
+        k_start_d = q_block_id * BS
+        offs_k_d = k_start_d + offs_d_arange
+
+        if ASSUME_FULL:
+            # No boundary branch — full unmasked diagonal load + causal mask.
+            k_bf16_d = tl.load(K_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+            v_bf16_d = tl.load(V_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+            s_d = tl.dot(q_bf16, tl.trans(k_bf16_d), out_dtype=tl.float32)
+            causal = offs_tok[:, None] >= offs_k_d[None, :]
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k_d[None, :]) < WINDOW
+                s_d = tl.where(causal & in_window, s_d, float("-inf"))
+            else:
+                s_d = tl.where(causal, s_d, float("-inf"))
+        else:
+            boundary = (q_start + BS) > seq_len
+            if boundary:
+                k_mask = offs_k_d < seq_len
+                k_bf16_d = tl.load(K_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :],
+                                   mask=k_mask[:, None], other=0.0)
+                v_bf16_d = tl.load(V_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :],
+                                   mask=k_mask[:, None], other=0.0)
+                s_d = tl.dot(q_bf16, tl.trans(k_bf16_d), out_dtype=tl.float32)
+                causal = offs_tok[:, None] >= offs_k_d[None, :]
+                allowed = causal & k_mask[None, :]
+                if WINDOW > 0:
+                    in_window = (offs_tok[:, None] - offs_k_d[None, :]) < WINDOW
+                    allowed = allowed & in_window
+                s_d = tl.where(allowed, s_d, float("-inf"))
+            else:
+                k_bf16_d = tl.load(K_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+                v_bf16_d = tl.load(V_ptr + offs_k_d[:, None] * stride_t + offs_d[None, :])
+                s_d = tl.dot(q_bf16, tl.trans(k_bf16_d), out_dtype=tl.float32)
+                causal = offs_tok[:, None] >= offs_k_d[None, :]
+                if WINDOW > 0:
+                    in_window = (offs_tok[:, None] - offs_k_d[None, :]) < WINDOW
+                    s_d = tl.where(causal & in_window, s_d, float("-inf"))
+                else:
+                    s_d = tl.where(causal, s_d, float("-inf"))
+
+        m_i = tl.max(s_d, axis=1)
+        p_d = tl.exp2(s_d - m_i[:, None])
+        l_i = tl.sum(p_d, axis=1)
+        acc = tl.dot(p_d.to(tl.bfloat16), v_bf16_d, out_dtype=tl.float32)
+
+    # Off-diag loop unchanged (no boundary semantics here in either mode).
+    if CONTIGUOUS_OFFDIAG:
+        k_block_id_start = tl.load(CI_ptr + ci_lo)
+        for ci in range(ci_lo, ci_hi - 1):
+            k_block_id = k_block_id_start + (ci - ci_lo)
+            k_start = k_block_id * BS
+            offs_k = k_start + offs_d_arange
+            k_bf16 = tl.load(K_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            v_bf16 = tl.load(V_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            s = tl.dot(q_bf16, tl.trans(k_bf16), out_dtype=tl.float32)
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k[None, :]) < WINDOW
+                s = tl.where(in_window, s, float("-inf"))
+            m_new = tl.maximum(m_i, tl.max(s, axis=1))
+            alpha = tl.exp2(m_i - m_new)
+            p = tl.exp2(s - m_new[:, None])
+            l_i = l_i * alpha + tl.sum(p, axis=1)
+            acc = tl.dot(p.to(tl.bfloat16), v_bf16, acc=acc * alpha[:, None], out_dtype=tl.float32)
+            m_i = m_new
+    else:
+        for ci in range(ci_lo, ci_hi - 1):
+            k_block_id = tl.load(CI_ptr + ci)
+            k_start = k_block_id * BS
+            offs_k = k_start + offs_d_arange
+            k_bf16 = tl.load(K_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            v_bf16 = tl.load(V_ptr + offs_k[:, None] * stride_t + offs_d[None, :])
+            s = tl.dot(q_bf16, tl.trans(k_bf16), out_dtype=tl.float32)
+            if WINDOW > 0:
+                in_window = (offs_tok[:, None] - offs_k[None, :]) < WINDOW
+                s = tl.where(in_window, s, float("-inf"))
+            m_new = tl.maximum(m_i, tl.max(s, axis=1))
+            alpha = tl.exp2(m_i - m_new)
+            p = tl.exp2(s - m_new[:, None])
+            l_i = l_i * alpha + tl.sum(p, axis=1)
+            acc = tl.dot(p.to(tl.bfloat16), v_bf16, acc=acc * alpha[:, None], out_dtype=tl.float32)
+            m_i = m_new
+
+    LN2: tl.constexpr = 0.6931471805599453
+    l_safe = tl.where(l_i > 0, l_i, 1.0)
+    if ASSUME_FULL:
+        o_tile = acc / l_safe[:, None]
+        lse_out = (m_i + tl.log2(l_safe)) * LN2
+    else:
+        q_mask = offs_tok < seq_len
+        o_tile = tl.where(q_mask[:, None], acc / l_safe[:, None], 0.0)
+        lse_out = tl.where(q_mask, (m_i + tl.log2(l_safe)) * LN2, float("-inf"))
+
+    O_ptr = O + bh_id * stride_h
+    tl.store(O_ptr + offs_tok[:, None] * stride_t + offs_d[None, :], o_tile.to(tl.bfloat16))
+    LSE_ptr = LSE + bh_id * stride_lh
+    tl.store(LSE_ptr + offs_tok, lse_out)
+
+
+_hydra_fwd_kernel = _kernel_decorator(_hydra_fwd_jit)
+
+
+def launch_attn_fwd(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    row_ptr: torch.Tensor,
+    col_idx: torch.Tensor,
+    seq_lens: torch.Tensor,
+    window: int = 0,
+    contiguous_offdiag: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if q.dim() != 4 or k.dim() != 4 or v.dim() != 4:
+        raise ValueError(f"q/k/v must be 4D; got {q.shape} {k.shape} {v.shape}")
+    batch_size, num_heads, t_max, head_dim = q.shape
+    num_kv_heads = k.shape[1]
+    if t_max % BLOCK_SIZE != 0:
+        raise ValueError(f"T ({t_max}) must be multiple of BLOCK_SIZE ({BLOCK_SIZE})")
+    if head_dim != HEAD_DIM:
+        raise ValueError(f"head_dim ({head_dim}) must equal HEAD_DIM ({HEAD_DIM})")
+
+    batch_heads = batch_size * num_heads
+    num_q_blocks = t_max // BLOCK_SIZE
+    o = torch.empty_like(q)
+    lse_3d = torch.empty((batch_size, num_heads, t_max), device=q.device, dtype=torch.float32)
+    lse_2d = lse_3d.view(batch_heads, t_max)
+    cih = col_idx.shape[2] if col_idx.ndim == 3 else col_idx.shape[1]
+    grid = (num_q_blocks * batch_heads,)
+
+    # ASSUME_FULL is safe iff EVERY seq_len equals t_max. Cheap GPU-resident
+    # check (no D2H sync if we just compare with a fused all-equal kernel,
+    # but for now use a tiny .item() — it's a single int and seq_lens lives
+    # in HBM near the launcher already).
+    if seq_lens.numel() == 1:
+        assume_full = int(seq_lens.item()) == t_max
+    else:
+        # Multi-batch: check min==max==t_max.
+        sl_min = int(seq_lens.min().item())
+        sl_max = int(seq_lens.max().item())
+        assume_full = (sl_min == t_max) and (sl_max == t_max)
+
+    common_kwargs = dict(
+        T_MAX=t_max, NUM_HEADS=num_heads, NUM_KV_HEADS=num_kv_heads,
+        SCALE=1.0 / math.sqrt(head_dim), BS=BLOCK_SIZE, D=head_dim,
+        WINDOW=int(window),
+        CONTIGUOUS_OFFDIAG=1 if contiguous_offdiag else 0,
+        ASSUME_FULL=1 if assume_full else 0,
+    )
+    if _DISABLE_AUTOTUNE:
+        common_kwargs["num_warps"] = _FWD_NUM_WARPS
+        common_kwargs["num_stages"] = _FWD_NUM_STAGES
+
+    _hydra_fwd_kernel[grid](
+        q, k, v, o, lse_2d,
+        row_ptr, col_idx, seq_lens, cih,
+        **common_kwargs,
+    )
+    return o, lse_3d

torch-ext/hydra/policy.py

@@ -0,0 +1,438 @@
+"""Runtime policy hooks for live attention throttling.
+
+This is intentionally small and deterministic. The policy can clamp an
+existing sliding-window request before CSR construction / kernel launch. It can
+also opt into converting dense causal attention to sliding-window attention,
+but that is disabled by default because it changes model semantics.
+"""
+from __future__ import annotations
+
+from collections import deque
+from dataclasses import asdict, dataclass
+import math
+import os
+from typing import Any, Mapping
+
+
+_MIB = 1024 * 1024
+
+
+def _parse_bool(value: object, default: bool = False) -> bool:
+    if value is None:
+        return default
+    if isinstance(value, bool):
+        return value
+    text = str(value).strip().lower()
+    if text in {"1", "true", "yes", "on", "enabled"}:
+        return True
+    if text in {"0", "false", "no", "off", "disabled"}:
+        return False
+    return default
+
+
+def _parse_int(value: object, default: int | None = None) -> int | None:
+    if value is None or value == "":
+        return default
+    return int(value)
+
+
+def _parse_int_tuple(value: object) -> tuple[int, ...] | None:
+    if value is None or value == "":
+        return None
+    if isinstance(value, str):
+        parts = [part.strip() for part in value.replace(";", ",").split(",")]
+        items = [int(part) for part in parts if part]
+    else:
+        items = [int(part) for part in value]  # type: ignore[arg-type]
+    if not items:
+        return None
+    if any(item <= 0 for item in items):
+        raise ValueError(f"layer windows must be positive integers, got {value!r}")
+    return tuple(items)
+
+
+def _parse_float(value: object, default: float) -> float:
+    if value is None or value == "":
+        return default
+    return float(value)
+
+
+def _parse_bytes(value: object, default: int) -> int:
+    if value is None or value == "":
+        return default
+    text = str(value).strip().lower()
+    scale = 1
+    for suffix, multiplier in (
+        ("gib", 1024**3),
+        ("gb", 1000**3),
+        ("mib", 1024**2),
+        ("mb", 1000**2),
+        ("kib", 1024),
+        ("kb", 1000),
+    ):
+        if text.endswith(suffix):
+            scale = multiplier
+            text = text[: -len(suffix)]
+            break
+    return int(float(text) * scale)
+
+
+@dataclass(frozen=True)
+class RuntimePolicy:
+    """Policy for request-time attention throttling."""
+
+    mode: str = "off"
+    max_window: int | None = None
+    min_window: int = 128
+    reserve_bytes: int = 512 * _MIB
+    utilization: float = 0.85
+    allow_dense_to_window: bool = False
+    fail_closed: bool = False
+    align_to_block: bool = True
+    layer_windows: tuple[int, ...] | None = None
+
+    @property
+    def enabled(self) -> bool:
+        return self.mode not in {"", "0", "off", "disabled", "none"}
+
+    @classmethod
+    def from_env(cls) -> "RuntimePolicy":
+        raw_mode = os.environ.get("HYDRA_POLICY", "off").strip().lower()
+        if raw_mode in {"1", "true", "yes", "on", "enabled"}:
+            raw_mode = "adaptive"
+        return cls(
+            mode=raw_mode,
+            max_window=_parse_int(os.environ.get("HYDRA_MAX_WINDOW")),
+            min_window=int(os.environ.get("HYDRA_MIN_WINDOW", "128")),
+            reserve_bytes=_parse_bytes(
+                os.environ.get("HYDRA_RESERVE_BYTES"),
+                int(float(os.environ.get("HYDRA_RESERVE_MB", "512")) * _MIB),
+            ),
+            utilization=_parse_float(os.environ.get("HYDRA_UTILIZATION"), 0.85),
+            allow_dense_to_window=_parse_bool(
+                os.environ.get("HYDRA_ALLOW_DENSE_THROTTLE"), False
+            ),
+            fail_closed=_parse_bool(os.environ.get("HYDRA_FAIL_CLOSED"), False),
+            align_to_block=_parse_bool(os.environ.get("HYDRA_ALIGN_WINDOW"), True),
+            layer_windows=_parse_int_tuple(
+                os.environ.get("HYDRA_LAYER_WINDOWS")
+            ),
+        )
+
+    @classmethod
+    def from_mapping(cls, data: Mapping[str, Any]) -> "RuntimePolicy":
+        return cls(
+            mode=str(data.get("mode", "adaptive")).strip().lower(),
+            max_window=_parse_int(data.get("max_window")),
+            min_window=int(data.get("min_window", 128)),
+            reserve_bytes=_parse_bytes(data.get("reserve_bytes"), 512 * _MIB),
+            utilization=float(data.get("utilization", 0.85)),
+            allow_dense_to_window=_parse_bool(data.get("allow_dense_to_window"), False),
+            fail_closed=_parse_bool(data.get("fail_closed"), False),
+            align_to_block=_parse_bool(data.get("align_to_block"), True),
+            layer_windows=_parse_int_tuple(data.get("layer_windows")),
+        )
+
+    def to_mapping(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+@dataclass(frozen=True)
+class PolicyDecision:
+    enabled: bool
+    action: str
+    reason: str
+    requested_window: int | None
+    effective_window: int | None
+    estimated_bytes: int | None = None
+    budget_bytes: int | None = None
+    free_bytes: int | None = None
+    total_bytes: int | None = None
+    layer_idx: int | None = None
+    layer_window: int | None = None
+
+    def to_mapping(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+_POLICY_OVERRIDE: RuntimePolicy | None = None
+_HISTORY: deque[PolicyDecision] = deque(maxlen=64)
+
+
+def set_runtime_policy(policy: RuntimePolicy | Mapping[str, Any] | None) -> None:
+    """Install a process-local policy override.
+
+    Passing ``None`` returns policy selection to environment variables.
+    """
+
+    global _POLICY_OVERRIDE
+    if policy is None:
+        _POLICY_OVERRIDE = None
+    elif isinstance(policy, RuntimePolicy):
+        _POLICY_OVERRIDE = policy
+    else:
+        _POLICY_OVERRIDE = RuntimePolicy.from_mapping(policy)
+
+
+def get_runtime_policy() -> RuntimePolicy:
+    return _POLICY_OVERRIDE or RuntimePolicy.from_env()
+
+
+def last_policy_decision() -> PolicyDecision | None:
+    return _HISTORY[-1] if _HISTORY else None
+
+
+def policy_history() -> list[dict[str, Any]]:
+    return [item.to_mapping() for item in _HISTORY]
+
+
+def _record(decision: PolicyDecision) -> PolicyDecision:
+    if decision.enabled:
+        _HISTORY.append(decision)
+    return decision
+
+
+def estimate_attention_bytes(
+    *,
+    batch_size: int,
+    query_heads: int,
+    query_tokens: int,
+    key_tokens: int,
+    head_dim: int,
+    dtype_bytes: int,
+    block_size: int,
+    sliding_window: int | None,
+    decode: bool = False,
+) -> int:
+    """Conservative estimate of extra bytes allocated by this attention call."""
+
+    output = batch_size * query_heads * query_tokens * head_dim * dtype_bytes
+    lse = batch_size * query_heads * query_tokens * 4
+    if decode:
+        return output + lse
+
+    q_blocks = math.ceil(query_tokens / block_size)
+    k_blocks = math.ceil(key_tokens / block_size)
+    row_ptr = batch_size * query_heads * (q_blocks + 1) * 4
+    seq_lens = batch_size * 4
+
+    if sliding_window is None or sliding_window >= key_tokens:
+        nnz_per_head = q_blocks * (q_blocks + 1) // 2
+    else:
+        blocks_per_row = min(k_blocks, math.ceil(sliding_window / block_size) + 1)
+        nnz_per_head = q_blocks * blocks_per_row
+    col_idx = batch_size * query_heads * nnz_per_head * 4
+    return output + lse + row_ptr + col_idx + seq_lens
+
+
+def _cuda_mem_info(device) -> tuple[int | None, int | None]:
+    if getattr(device, "type", None) != "cuda":
+        return None, None
+    try:
+        import torch
+
+        free, total = torch.cuda.mem_get_info(device)
+        return int(free), int(total)
+    except Exception:
+        return None, None
+
+
+def _budget(policy: RuntimePolicy, free_bytes: int | None) -> int | None:
+    if free_bytes is None:
+        return None
+    usable = max(0, free_bytes - policy.reserve_bytes)
+    return int(usable * policy.utilization)
+
+
+def _align_window(window: int, block_size: int, policy: RuntimePolicy) -> int:
+    window = max(policy.min_window, window)
+    if policy.align_to_block:
+        window = max(block_size, (window // block_size) * block_size)
+    return window
+
+
+def _layer_window_for(
+    policy: RuntimePolicy, layer_idx: int | None, block_size: int
+) -> int | None:
+    if layer_idx is None or not policy.layer_windows:
+        return None
+    if layer_idx < 0:
+        return None
+    raw_window = policy.layer_windows[layer_idx % len(policy.layer_windows)]
+    return _align_window(raw_window, block_size, policy)
+
+
+def _fit_window_for_budget(
+    *,
+    budget_bytes: int,
+    fixed_bytes: int,
+    batch_size: int,
+    query_heads: int,
+    query_tokens: int,
+    block_size: int,
+    policy: RuntimePolicy,
+) -> int:
+    q_blocks = math.ceil(query_tokens / block_size)
+    bytes_per_block_per_row = batch_size * query_heads * q_blocks * 4
+    if bytes_per_block_per_row <= 0:
+        return policy.min_window
+    remaining = max(0, budget_bytes - fixed_bytes)
+    # ``fixed_bytes`` is estimated with a one-token window, which still maps to
+    # two CSR blocks per row: one reach block plus the diagonal block. Add any
+    # remaining block budget to that floor, then subtract the diagonal to get
+    # token-window reach.
+    blocks_per_row = 2 + max(0, remaining // bytes_per_block_per_row)
+    # CSR includes the diagonal block in addition to window reach.
+    window_blocks = max(1, int(blocks_per_row) - 1)
+    return _align_window(window_blocks * block_size, block_size, policy)
+
+
+def apply_runtime_policy(
+    q,
+    k,
+    v,
+    *,
+    sliding_window: int | None,
+    block_size: int,
+    head_dim: int,
+    layer_idx: int | None = None,
+) -> PolicyDecision:
+    """Apply the active policy and return the effective sliding window."""
+
+    policy = get_runtime_policy()
+    requested = sliding_window
+    if not policy.enabled:
+        return PolicyDecision(False, "off", "policy disabled", requested, requested)
+
+    batch_size, query_heads, query_tokens, dim = q.shape
+    key_tokens = k.shape[2]
+    decode = query_tokens == 1
+    dtype_bytes = q.element_size()
+    free, total = _cuda_mem_info(q.device)
+    budget = _budget(policy, free)
+    effective = requested
+    layer_window = _layer_window_for(policy, layer_idx, block_size)
+    action = "pass"
+    reason = "within policy"
+
+    if policy.max_window is not None:
+        max_window = _align_window(policy.max_window, block_size, policy)
+        if effective is None and policy.allow_dense_to_window:
+            effective = max_window
+            action = "dense_to_window"
+            reason = f"dense attention converted to sw={effective}"
+        elif effective is not None and effective > max_window:
+            effective = max_window
+            action = "clamp_max_window"
+            reason = f"sliding_window clamped to max_window={effective}"
+
+    if layer_window is not None:
+        if effective is None and policy.allow_dense_to_window:
+            effective = layer_window
+            action = "layer_window"
+            reason = f"dense attention converted to layer_window={effective}"
+        elif effective is not None and effective > layer_window:
+            effective = layer_window
+            action = "layer_window"
+            reason = f"sliding_window clamped to layer_window={effective}"
+
+    estimate = estimate_attention_bytes(
+        batch_size=batch_size,
+        query_heads=query_heads,
+        query_tokens=query_tokens,
+        key_tokens=key_tokens,
+        head_dim=dim,
+        dtype_bytes=dtype_bytes,
+        block_size=block_size,
+        sliding_window=effective,
+        decode=decode,
+    )
+
+    if policy.mode == "adaptive" and budget is not None and estimate > budget and not decode:
+        fixed = estimate_attention_bytes(
+            batch_size=batch_size,
+            query_heads=query_heads,
+            query_tokens=query_tokens,
+            key_tokens=key_tokens,
+            head_dim=dim,
+            dtype_bytes=dtype_bytes,
+            block_size=block_size,
+            sliding_window=1,
+            decode=False,
+        )
+        if effective is None and policy.allow_dense_to_window:
+            effective = policy.max_window or key_tokens
+        if effective is not None:
+            fit_window = _fit_window_for_budget(
+                budget_bytes=budget,
+                fixed_bytes=fixed,
+                batch_size=batch_size,
+                query_heads=query_heads,
+                query_tokens=query_tokens,
+                block_size=block_size,
+                policy=policy,
+            )
+            if fit_window < effective:
+                effective = fit_window
+                action = "budget_throttle"
+                reason = f"estimated attention bytes exceeded budget; sw={effective}"
+                estimate = estimate_attention_bytes(
+                    batch_size=batch_size,
+                    query_heads=query_heads,
+                    query_tokens=query_tokens,
+                    key_tokens=key_tokens,
+                    head_dim=dim,
+                    dtype_bytes=dtype_bytes,
+                    block_size=block_size,
+                    sliding_window=effective,
+                    decode=False,
+                )
+        elif policy.fail_closed:
+            decision = PolicyDecision(
+                True,
+                "reject",
+                "dense attention exceeds policy budget and dense throttling is disabled",
+                requested,
+                effective,
+                estimate,
+                budget,
+                free,
+                total,
+                layer_idx,
+                layer_window,
+            )
+            _record(decision)
+            raise MemoryError(decision.reason)
+
+    if policy.fail_closed and budget is not None and estimate > budget:
+        decision = PolicyDecision(
+            True,
+            "reject",
+            "attention estimate still exceeds policy budget after throttling",
+            requested,
+            effective,
+            estimate,
+            budget,
+            free,
+            total,
+            layer_idx,
+            layer_window,
+        )
+        _record(decision)
+        raise MemoryError(decision.reason)
+
+    return _record(
+        PolicyDecision(
+            True,
+            action,
+            reason,
+            requested,
+            effective,
+            estimate,
+            budget,
+            free,
+            total,
+            layer_idx,
+            layer_window,
+        )
+    )

torch-ext/hydra/sw_sinks_csr.py

@@ -0,0 +1,93 @@
+"""Sliding-window + attention-sinks CSR builder.
+
+StreamingLLM (Xiao et al. 2023, arXiv:2309.17453) showed that always keeping
+the first N tokens ("attention sinks") visible recovers semantic quality
+when a non-SW-trained LM is forced into a sliding-window attention regime.
+Without sinks, the softmax distribution loses the anchor it implicitly
+deposits on the first tokens during pretraining; with sinks, models like
+Qwen2.5-Coder / Qwen3 stay coherent at long context under SW=4096.
+
+This builder produces the same ``(row_ptr, col_idx, seq_lens)`` triple as
+``build_sliding_window_csr`` but for each Q-block additionally includes the
+first ``ceil(sinks / BLOCK_K)`` K-blocks. Deduplicated when the SW range
+already covers the sinks. Per-row col_idx is sorted ascending with the
+diagonal block as the last entry (kernel convention — see ``csr.py``).
+"""
+from __future__ import annotations
+
+import math
+
+import torch
+
+from .csr import _broadcast_csr
+
+
+def build_sw_sinks_csr(
+    window: int,
+    seq_len: int,
+    block_size: int,
+    batch_size: int,
+    num_heads: int,
+    device: torch.device | str,
+    sinks: int = 4,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Sliding-window + first-N-sinks causal CSR.
+
+    Same shape contract as ``build_sliding_window_csr``: returns
+    ``(row_ptr, col_idx, seq_lens)`` broadcast across (B, H).
+
+    For Q-block ``qb``:
+      - Off-diagonal SW range (block ids): ``[k_min_sw, qb)`` where
+        ``k_min_sw = max(0, (qb*BS - window + 1)) // BS`` if positive
+        else 0.
+      - PLUS sink K-blocks: ``[0, sinks_blocks)`` with
+        ``sinks_blocks = ceil(sinks / BS)``.
+      - PLUS diagonal placeholder ``qb`` (always the last entry).
+
+    The union is deduplicated and sorted ascending; the diagonal is
+    appended last per the kernel's row-walker contract.
+
+    Edge cases:
+      - sinks <= 0 OR sinks_blocks == 0: degrades to plain SW.
+      - sinks_blocks >= num_q_blocks: degrades to dense causal.
+      - window covers the sinks (qb*BS < window): sinks already
+        included; dedup makes this a no-op.
+    """
+    if window <= 0:
+        raise ValueError(f"window must be positive, got {window}")
+    if block_size <= 0:
+        raise ValueError(f"block_size must be positive, got {block_size}")
+    if sinks < 0:
+        raise ValueError(f"sinks must be non-negative, got {sinks}")
+
+    num_q_blocks = math.ceil(seq_len / block_size)
+    sinks_blocks = math.ceil(sinks / block_size) if sinks > 0 else 0
+    # Clip so we never claim a sink block past the sequence end.
+    sinks_blocks = min(sinks_blocks, num_q_blocks)
+
+    row_ptr_row = [0]
+    col_idx_row: list[int] = []
+    for q_block in range(num_q_blocks):
+        # SW left-edge in K-block space (same logic as build_sliding_window_csr).
+        left_bound = q_block * block_size - window + 1
+        if left_bound <= 0:
+            k_min_sw = 0
+        else:
+            k_min_sw = left_bound // block_size
+
+        sw_range_end = q_block  # exclusive (diagonal excluded)
+        sw_lo = min(k_min_sw, sw_range_end)
+
+        # Sink block ids that fall strictly before the diagonal AND strictly
+        # before the SW range start (anything inside SW is already covered).
+        sink_lo = 0
+        sink_hi = min(sinks_blocks, sw_lo)  # exclusive
+
+        if sink_hi > sink_lo:
+            col_idx_row.extend(range(sink_lo, sink_hi))
+        col_idx_row.extend(range(sw_lo, sw_range_end))
+        # Diagonal placeholder must be the last entry per kernel contract.
+        col_idx_row.append(q_block)
+        row_ptr_row.append(len(col_idx_row))
+
+    return _broadcast_csr(row_ptr_row, col_idx_row, batch_size, num_heads, seq_len, device)