flashrt-structures 0.2.0__py3-none-any.whl

flashrt_structures/impls/attention_core/fa4_cute.py

@@ -0,0 +1,154 @@
+"""attention_core — the FlashAttention-4 (CuTe DSL) dense form.
+
+The community ``kernels-community/flash-attn4`` package, reused as-is:
+the measured Thor comparison put it at operational parity with the
+native FA4 backend across the GROOT profiles (D48 / D72 / causal GQA
+D128), so no FlashRT twin is published. Numerics-preserving BF16 —
+in the variant family's precision ordering it sits directly after FA2,
+ahead of the FP8 forms.
+
+Same stateless dense seam as the FA2 dense form: complete Q/K/V per
+call, host SDPA layout, allowed-ranges packing. GQA is native
+(``pack_gqa``); custom softmax scales pass straight through. The DSL
+JIT compiles per shape at first call, which the bind smoke absorbs at
+bind time — a device or dependency the DSL cannot serve surfaces as a
+bind refusal there, not inside the host's forward.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardedSeam
+from .fa2_seqused import _allowed_ranges
+
+KERNEL_DEP = {
+    "provider": "huggingface_kernels",
+    "repo": "kernels-community/flash-attn4",
+    "version": ">=0",
+}
+
+
+class DenseAttentionFa4Cute(GuardedSeam, torch.nn.Module):
+    """FA4 (CuTe DSL) replacement for an ordinary dense SDPA call.
+
+    Inputs and outputs use the host SDPA layout ``[B, H, S, D]``; the
+    kernel consumes ``[B, S, H, D]``.
+    """
+
+    def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device,
+                 allowed_ranges=None):
+        super().__init__()
+        b, heads, seq_q, head_dim = q_shape
+        kb, kv_heads, seq_kv, kv_dim = kv_shape
+        if kb != b or kv_dim != head_dim:
+            raise ValueError(
+                "attention_core fa4_cute: Q and KV batch/head "
+                "dimensions differ")
+        if heads % kv_heads:
+            raise ValueError(
+                "attention_core fa4_cute: query heads must be "
+                "divisible by KV heads")
+        self.q_shape = tuple(q_shape)
+        self.kv_shape = tuple(kv_shape)
+        self.allowed_ranges = tuple(allowed_ranges or ())
+        kern = hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+        self._fn = kern.flash_attn_func
+        if self.allowed_ranges:
+            packed_seq = sum(hi - lo for lo, hi in self.allowed_ranges)
+            self.register_buffer("packed_k", torch.empty(
+                b, packed_seq, kv_heads, head_dim, device=device,
+                dtype=dtype))
+            self.register_buffer("packed_v",
+                                 torch.empty_like(self.packed_k))
+        self._frt_arm(
+            dtypes=(dtype,), device=torch.device(device),
+            k=int(head_dim), rows=int(b * heads * seq_q))
+
+    def forward(self, query, key, value, *, scale=None):
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        if tuple(query.shape) != self.q_shape:
+            raise ValueError(
+                "attention_core fa4_cute: query shape moved from "
+                f"{self.q_shape} to {tuple(query.shape)}")
+        if tuple(key.shape) != self.kv_shape or value.shape != key.shape:
+            raise ValueError(
+                "attention_core fa4_cute: key/value shape moved from "
+                f"{self.kv_shape} to {tuple(key.shape)}/"
+                f"{tuple(value.shape)}")
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        if not q.is_contiguous():
+            q = q.contiguous()
+        if self.allowed_ranges:
+            offset = 0
+            for lo, hi in self.allowed_ranges:
+                length = hi - lo
+                self.packed_k[:, offset:offset + length].copy_(k[:, lo:hi])
+                self.packed_v[:, offset:offset + length].copy_(v[:, lo:hi])
+                offset += length
+            k, v = self.packed_k, self.packed_v
+        else:
+            if not k.is_contiguous():
+                k = k.contiguous()
+            if not v.is_contiguous():
+                v = v.contiguous()
+        out = self._fn(q, k, v, softmax_scale=scale, causal=False)
+        if isinstance(out, tuple):
+            out = out[0]
+        return out.transpose(1, 2).to(query.dtype)
+
+
+def bind_dense_attention(captures):
+    """Bind one stateless dense FA4-CuTe core from repeated captures.
+
+    Same qualification walk as the family's other dense binders. The
+    DSL owns its own shape envelope, so there is no head-dim table
+    here: the bind smoke below runs the real entry at the captured
+    shape, and what the DSL cannot compile is a bind refusal.
+    """
+    if not captures:
+        raise ValueError("attention_core fa4_cute: no captures")
+    first = captures[0]
+    query, key, value = first["q"], first["key"], first["value"]
+    allowed_ranges = _allowed_ranges(first.get("mask"))
+    if allowed_ranges is None:
+        return None
+    expected = (tuple(query.shape), tuple(key.shape), tuple(value.shape),
+                query.dtype, key.dtype, value.dtype)
+    for capture in captures[1:]:
+        got = (
+            tuple(capture["q"].shape),
+            tuple(capture["key"].shape),
+            tuple(capture["value"].shape),
+            capture["q"].dtype,
+            capture["key"].dtype,
+            capture["value"].dtype,
+        )
+        if got != expected:
+            raise ValueError(
+                "attention_core fa4_cute: shape, dtype, or mask moved "
+                f"within one calibration call: {expected} -> {got}")
+        if _allowed_ranges(capture.get("mask")) != allowed_ranges:
+            raise ValueError(
+                "attention_core fa4_cute: mask pattern moved within "
+                "one calibration call")
+    if not (query.dtype == key.dtype == value.dtype):
+        raise ValueError("attention_core fa4_cute: Q/K/V dtypes differ")
+    bound = DenseAttentionFa4Cute(
+        query.shape, key.shape, query.dtype, query.device,
+        allowed_ranges=allowed_ranges)
+    with torch.no_grad():
+        probe = bound(torch.zeros_like(query),
+                      torch.zeros_like(key),
+                      torch.zeros_like(value))
+    if probe.shape != query.shape or not torch.isfinite(probe).all():
+        raise ValueError(
+            "attention_core fa4_cute: bind smoke produced shape "
+            f"{tuple(probe.shape)}, "
+            f"finite={bool(torch.isfinite(probe).all())}")
+    return bound

flashrt_structures/impls/attention_core/fa4_fp8.py

@@ -0,0 +1,178 @@
+"""attention_core — the FA4 (Blackwell FP8) dense form.
+
+The SM100-family twin of :class:`.fa2_seqused.DenseAttention`: the same
+stateless dense seam — complete Q/K/V every call, host SDPA layout,
+allowed-ranges packing — executed by the
+``flashrt/fp8-cross-attention-blackwell`` kernel: non-causal FP8 GQA
+attention at head_dim 128, BF16 out. The two variants split the
+hardware between them through their packages' own arch declarations
+(this one ships ``10.0a/11.0a`` builds, the FA2 runtime ships none for
+those majors), so selection is the ordinary refusal machinery rather
+than a second table: the family binder tries FA2 first and falls to
+this form where FA2's kernel refuses the device. On the devices this
+form serves, FP8 attention is the production hot path, and the parity
+gates downstream judge its quantization like any other impl's.
+
+Activation scales are per-tensor static, calibrated from the same real
+captures the qualification reads (amax over every capture, house
+formula). The kernel bakes the ``1/sqrt(head_dim)`` softmax convention;
+a host that calls with any other scale is refused at that call.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardedSeam
+from .fa2_seqused import _allowed_ranges
+
+KERNEL_DEP = {
+    "provider": "huggingface_kernels",
+    "repo": "flashrt/fp8-cross-attention-blackwell",
+    "version": ">=1",
+}
+
+_FP8 = torch.float8_e4m3fn
+_FP8_MAX = 448.0
+_HEAD_DIM = 128  # the kernel's contract, exact
+
+
+class DenseAttentionFa4(GuardedSeam, torch.nn.Module):
+    """FA4 replacement for an ordinary dense SDPA call.
+
+    Inputs and outputs use the host SDPA layout ``[B, H, S, D]``; the
+    kernel consumes ``[B, S, H, D]``. Q/K/V are quantized per call with
+    the calibrated static scales — elementwise work a compiled or
+    packaged graph fuses into its neighbours.
+    """
+
+    def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device,
+                 scales: tuple[float, float, float],
+                 allowed_ranges=None):
+        super().__init__()
+        b, heads, seq_q, head_dim = q_shape
+        kb, kv_heads, seq_kv, kv_dim = kv_shape
+        if kb != b or kv_dim != head_dim:
+            raise ValueError(
+                "attention_core fa4: Q and KV batch/head dimensions differ")
+        if heads % kv_heads:
+            raise ValueError(
+                "attention_core fa4: query heads must be divisible by "
+                "KV heads")
+        if head_dim != _HEAD_DIM:
+            raise ValueError(
+                f"attention_core fa4: head_dim {head_dim} outside the "
+                f"kernel contract ({_HEAD_DIM})")
+        self.q_shape = tuple(q_shape)
+        self.kv_shape = tuple(kv_shape)
+        self.allowed_ranges = tuple(allowed_ranges or ())
+        self._kfa = hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+        self._fn = self._kfa.fp8_gqa_cross_attention_bf16
+        self._qs, self._ks, self._vs = (float(s) for s in scales)
+        if self.allowed_ranges:
+            packed_seq = sum(hi - lo for lo, hi in self.allowed_ranges)
+            self.register_buffer("packed_k", torch.empty(
+                b, packed_seq, kv_heads, head_dim, device=device,
+                dtype=dtype))
+            self.register_buffer("packed_v",
+                                 torch.empty_like(self.packed_k))
+        self._frt_arm(
+            dtypes=(dtype,), device=torch.device(device),
+            k=int(head_dim), rows=int(b * heads * seq_q))
+
+    def _quant(self, t: torch.Tensor, scale: float) -> torch.Tensor:
+        return (t.float() / scale).clamp(-_FP8_MAX, _FP8_MAX).to(_FP8)
+
+    def forward(self, query, key, value, *, scale=None):
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        if tuple(query.shape) != self.q_shape:
+            raise ValueError(
+                "attention_core fa4: query shape moved from "
+                f"{self.q_shape} to {tuple(query.shape)}")
+        if tuple(key.shape) != self.kv_shape or value.shape != key.shape:
+            raise ValueError(
+                "attention_core fa4: key/value shape moved from "
+                f"{self.kv_shape} to {tuple(key.shape)}/"
+                f"{tuple(value.shape)}")
+        if scale is not None and abs(
+                float(scale) - self.q_shape[-1] ** -0.5) > 1e-9:
+            raise ValueError(
+                "attention_core fa4: the kernel bakes the default "
+                "softmax scale; this host passes another")
+        q = query.transpose(1, 2)
+        k = key.transpose(1, 2)
+        v = value.transpose(1, 2)
+        if self.allowed_ranges:
+            offset = 0
+            for lo, hi in self.allowed_ranges:
+                length = hi - lo
+                self.packed_k[:, offset:offset + length].copy_(k[:, lo:hi])
+                self.packed_v[:, offset:offset + length].copy_(v[:, lo:hi])
+                offset += length
+            k, v = self.packed_k, self.packed_v
+        out = self._fn(
+            self._quant(q, self._qs).contiguous(),
+            self._quant(k, self._ks).contiguous(),
+            self._quant(v, self._vs).contiguous(),
+            query_scale=self._qs, key_scale=self._ks,
+            value_scale=self._vs)
+        return out.transpose(1, 2).to(query.dtype)
+
+
+def _amax_over(captures, key) -> float:
+    amax = 0.0
+    for cap in captures:
+        amax = max(amax, float(cap[key].detach().float().abs().max()))
+    return max(amax / _FP8_MAX, 1e-6)
+
+
+def bind_dense_attention(captures):
+    """Bind one stateless dense FA4 core from repeated host captures.
+
+    Same qualification walk as the FA2 dense binder — stable shapes and
+    dtypes across captures, a mask expressible as contiguous allowed
+    ranges — plus the kernel's own contract (head_dim 128) and the
+    per-tensor scale calibration this form adds. Returns ``None`` when
+    the shape qualification fails so the caller can keep its path;
+    raises when the kernel package refuses the device, so the family
+    binder can record it and move on.
+    """
+    if not captures:
+        raise ValueError("attention_core fa4: no captures")
+    first = captures[0]
+    query, key, value = first["q"], first["key"], first["value"]
+    if query.shape[-1] != _HEAD_DIM:
+        return None
+    allowed_ranges = _allowed_ranges(first.get("mask"))
+    if allowed_ranges is None:
+        return None
+    expected = (tuple(query.shape), tuple(key.shape), tuple(value.shape),
+                query.dtype, key.dtype, value.dtype)
+    for capture in captures[1:]:
+        got = (
+            tuple(capture["q"].shape),
+            tuple(capture["key"].shape),
+            tuple(capture["value"].shape),
+            capture["q"].dtype,
+            capture["key"].dtype,
+            capture["value"].dtype,
+        )
+        if got != expected:
+            raise ValueError(
+                "attention_core fa4: shape, dtype, or mask moved within "
+                f"one calibration call: {expected} -> {got}")
+        if _allowed_ranges(capture.get("mask")) != allowed_ranges:
+            raise ValueError(
+                "attention_core fa4: mask pattern moved within one "
+                "calibration call")
+    if not (query.dtype == key.dtype == value.dtype):
+        raise ValueError("attention_core fa4: Q/K/V dtypes differ")
+    scales = (_amax_over(captures, "q"),
+              _amax_over(captures, "key"),
+              _amax_over(captures, "value"))
+    return DenseAttentionFa4(
+        query.shape, key.shape, query.dtype, query.device,
+        scales, allowed_ranges=allowed_ranges)

flashrt_structures/impls/attention_core/masked_mha.py

@@ -0,0 +1,158 @@
+"""attention_core — the allocation-free masked-MHA dense form.
+
+The ``flashrt/masked-mha-runtime`` package: a padded-length MHA whose
+softmax reads and writes only the valid key columns, with every buffer
+caller-owned — the form that removed the per-call logits pre-fill
+sweep and, with it, the graph-replay nondeterminism of reading
+uninitialized padding (replays are bitwise). BF16/FP16 at the host's
+own precision; in the family's ordering it follows the FA4 forms and
+precedes the FP8 one.
+
+The kernel's layout is batch-free ``(S, H, D)`` per tensor, so this
+form serves the batch-of-one sites the packed dense seam sees
+everywhere in the VLA hosts; a batched site stays with the other
+variants. Masks reduce to contiguous allowed ranges exactly as in the
+FA2 dense form — packed KV rows are the valid length, which is the
+whole masking contract.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from .. import hub_kernel
+from ...guard import PROCEED, GuardedSeam
+from .fa2_seqused import _allowed_ranges
+
+KERNEL_DEP = {
+    "provider": "huggingface_kernels",
+    "repo": "flashrt/masked-mha-runtime",
+    "version": ">=1",
+}
+
+
+class DenseAttentionMaskedMha(GuardedSeam, torch.nn.Module):
+    """Masked-MHA replacement for a batch-of-one dense SDPA call.
+
+    Host layout ``[1, H, S, D]`` in and out; the kernel consumes
+    ``(S, H, D)`` with a caller-owned padded logits scratch and output,
+    both allocated once here — the hot path allocates nothing.
+    """
+
+    def __init__(self, q_shape, kv_shape, dtype: torch.dtype, device,
+                 allowed_ranges=None):
+        super().__init__()
+        b, heads, seq_q, head_dim = q_shape
+        kb, kv_heads, seq_kv, kv_dim = kv_shape
+        if b != 1 or kb != 1:
+            raise ValueError(
+                "attention_core masked_mha: the kernel layout is "
+                "batch-free; only batch-of-one sites qualify")
+        if kv_dim != head_dim:
+            raise ValueError(
+                "attention_core masked_mha: Q and KV head dims differ")
+        if heads != kv_heads:
+            raise ValueError(
+                "attention_core masked_mha: MHA form; GQA sites take "
+                "the FA4 variants")
+        self.q_shape = tuple(q_shape)
+        self.kv_shape = tuple(kv_shape)
+        self.allowed_ranges = tuple(allowed_ranges or ())
+        kern = hub_kernel(KERNEL_DEP["repo"], KERNEL_DEP["version"])
+        self._fn = kern.forward_static
+        packed_seq = (sum(hi - lo for lo, hi in self.allowed_ranges)
+                      if self.allowed_ranges else seq_kv)
+        self._packed_seq = packed_seq
+        if self.allowed_ranges:
+            self.register_buffer("packed_k", torch.empty(
+                packed_seq, kv_heads, head_dim, device=device,
+                dtype=dtype))
+            self.register_buffer("packed_v",
+                                 torch.empty_like(self.packed_k))
+        stride = (packed_seq + 7) // 8 * 8
+        self.register_buffer("_logits", torch.empty(
+            heads, seq_q, stride, device=device, dtype=dtype))
+        self.register_buffer("_out", torch.empty(
+            seq_q, heads, head_dim, device=device, dtype=dtype))
+        self._frt_arm(
+            dtypes=(dtype,), device=torch.device(device),
+            k=int(head_dim), rows=int(heads * seq_q))
+
+    def forward(self, query, key, value, *, scale=None):
+        admitted = self._frt_admit(query)
+        if admitted is not PROCEED:
+            return admitted
+        if tuple(query.shape) != self.q_shape:
+            raise ValueError(
+                "attention_core masked_mha: query shape moved from "
+                f"{self.q_shape} to {tuple(query.shape)}")
+        if tuple(key.shape) != self.kv_shape or value.shape != key.shape:
+            raise ValueError(
+                "attention_core masked_mha: key/value shape moved from "
+                f"{self.kv_shape} to {tuple(key.shape)}/"
+                f"{tuple(value.shape)}")
+        q = query[0].transpose(0, 1).contiguous()      # (S_q, H, D)
+        k = key[0].transpose(0, 1)
+        v = value[0].transpose(0, 1)
+        if self.allowed_ranges:
+            offset = 0
+            for lo, hi in self.allowed_ranges:
+                length = hi - lo
+                self.packed_k[offset:offset + length].copy_(k[lo:hi])
+                self.packed_v[offset:offset + length].copy_(v[lo:hi])
+                offset += length
+            k, v = self.packed_k, self.packed_v
+        else:
+            k = k.contiguous()
+            v = v.contiguous()
+        out = self._fn(q, k, v, logits=self._logits, out=self._out,
+                       scale=scale)
+        return out.transpose(0, 1).unsqueeze(0).to(query.dtype)
+
+
+def bind_dense_attention(captures):
+    """Bind one masked-MHA core from repeated host captures."""
+    if not captures:
+        raise ValueError("attention_core masked_mha: no captures")
+    first = captures[0]
+    query, key, value = first["q"], first["key"], first["value"]
+    if query.shape[0] != 1 or query.shape[1] != key.shape[1]:
+        return None
+    allowed_ranges = _allowed_ranges(first.get("mask"))
+    if allowed_ranges is None:
+        return None
+    expected = (tuple(query.shape), tuple(key.shape), tuple(value.shape),
+                query.dtype, key.dtype, value.dtype)
+    for capture in captures[1:]:
+        got = (
+            tuple(capture["q"].shape),
+            tuple(capture["key"].shape),
+            tuple(capture["value"].shape),
+            capture["q"].dtype,
+            capture["key"].dtype,
+            capture["value"].dtype,
+        )
+        if got != expected:
+            raise ValueError(
+                "attention_core masked_mha: shape, dtype, or mask "
+                f"moved within one calibration call: {expected} -> {got}")
+        if _allowed_ranges(capture.get("mask")) != allowed_ranges:
+            raise ValueError(
+                "attention_core masked_mha: mask pattern moved within "
+                "one calibration call")
+    if not (query.dtype == key.dtype == value.dtype):
+        raise ValueError(
+            "attention_core masked_mha: Q/K/V dtypes differ")
+    bound = DenseAttentionMaskedMha(
+        query.shape, key.shape, query.dtype, query.device,
+        allowed_ranges=allowed_ranges)
+    with torch.no_grad():
+        probe = bound(torch.zeros_like(query),
+                      torch.zeros_like(key),
+                      torch.zeros_like(value))
+    if probe.shape != query.shape or not torch.isfinite(probe).all():
+        raise ValueError(
+            "attention_core masked_mha: bind smoke produced shape "
+            f"{tuple(probe.shape)}, "
+            f"finite={bool(torch.isfinite(probe).all())}")
+    return bound