[Cake 8/13] gqa: packed causal GQA attention, head dim 256

miles_ext_kernels/cake/_impl/gqa/__init__.py

@@ -0,0 +1,18 @@
+"""
+Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+Licensed under the Apache License, Version 2.0.
+https://www.apache.org/licenses/LICENSE-2.0
+
+The implementation of ``miles_ext_kernels.cake.gqa`` on its vendored package ``cake_gqa_train`` (:mod:`.ops`),
+imported on first use of an entry point. Importing it imports torch.
+"""
+
+from miles_ext_kernels.cake._runtime import loader
+
+loader.install_op("gqa")  # the runtime loader on the op's package, before it can build
+
+from miles_ext_kernels.cake.gqa import check_supported, is_supported  # noqa: E402
+
+from .ops import prebuild, release, varlen_gqa_attention  # noqa: E402
+
+__all__ = ["check_supported", "is_supported", "prebuild", "release", "varlen_gqa_attention"]

miles_ext_kernels/cake/_impl/gqa/ops.py

@@ -0,0 +1,198 @@
+"""
+Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+Licensed under the Apache License, Version 2.0.
+https://www.apache.org/licenses/LICENSE-2.0
+
+Packed causal GQA attention, head dimension 256, forward and backward on the vendored ``cake_gqa_train`` package
+(FlashInfer's ``flashinfer.gqa_train`` backend), SM100a / SM103a / SM107a.
+
+Semantics. ``q [T_q, H_q, 256]``, ``k`` / ``v [T_k, H_kv, 256]`` (BF16, ``H_q / H_kv`` even), segments packed by
+int32 device ``cu_seqlens_q`` / ``cu_seqlens_k`` (``cu[0] = 0``, non-decreasing, ``cu_q[B] = T_q``,
+``cu_k[B] = T_k``). Query row ``i`` of a segment with ``q_len`` rows and ``k_len`` keys sees the keys
+``j <= i + k_len - q_len`` (FlashAttention's bottom-right alignment); a row that sees no key gets a zero output and a
+zero gradient.
+
+``max_seqlen_k``. The backward main kernel runs one CTA pair per 128 keys of ``max_seqlen_k``, per key/value head and
+segment, and a pair past its segment's last key block retires at once. A value below the longest segment's key count
+leaves key blocks out: their dK / dV are never written and their dQ terms are missing, without an error. A larger value
+is correct but launches pairs that only retire, which slows the backward down the more segments a call packs (``T_k``
+is always correct, and the slowest). The kernels never read the boundaries on the host, so only the caller keeps this
+contract; the option ``cake_gqa_check_seqlens`` (``--cake-gqa-check-seqlens``) checks it at every call, at the price
+of one device synchronization.
+
+Numerics. FP32 accumulation throughout, P and dS rounded to BF16 only as MMA operands. The output, the LSE and the BF16
+dK / dV repeat bit for bit run to run, whatever valid ``max_seqlen_k`` is given; dQ is summed in FP32 with atomic adds,
+one partial per 128-key block in a scheduling order, and rounded to BF16 once, so it is not reproducible bit for bit.
+
+Memory. Each backward takes its scratch (FP32 ``[H_q / 2, rows_pad, 256]`` and two ``[H_q / 2, rows_pad]``,
+``rows_pad = round_up(2 T_q + 128 B, 128)``) from the runtime's workspace pool on the current stream, and each forward
+zeroes a fresh ``uint32`` tile counter; autograd saves ``(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)`` and scalars
+only. Every launch is a plain launch on the current stream, and nothing synchronizes the device.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from miles_ext_kernels.cake._runtime import config, loader, workspace
+from miles_ext_kernels.cake._vendor.cake_gqa_train import cake_backend, cake_jit
+
+__all__ = [
+    "HEAD_DIM",
+    "PACKAGE",
+    "STAGES",
+    "PoolWorkspace",
+    "check_seqlens",
+    "libraries",
+    "prebuild",
+    "release",
+    "varlen_gqa_attention",
+]
+
+OP = "gqa"
+PACKAGE = "cake_gqa_train"
+HEAD_DIM = cake_backend.HEAD_DIM
+# The stages one forward and backward of varlen_gqa_attention launch (BF16 dK / dV), in launch order.
+STAGES = ("fwd", "bwd_pre", "bwd_bf16", "bwd_post")
+
+
+def libraries(arch: str) -> list[str]:
+    """The ``build.json`` libraries of :data:`STAGES` on ``arch``."""
+    return [cake_backend.library_name(cake_backend.stage_record(stage, arch), arch) for stage in STAGES]
+
+
+def prebuild(*, device: Any = None, max_workers: int | None = None) -> list[Path]:
+    """Build the libraries of :data:`STAGES` for the architecture of ``device`` (default: the current one).
+
+    No model parameter: every head layout, length and option runs these four libraries.
+    """
+    return loader.prebuild_libraries(PACKAGE, libraries(loader.device_arch(device)), max_workers=max_workers)
+
+
+def release() -> None:
+    """Nothing to drop here: the backward scratch lives in the runtime's workspace pool, the tile counter per call."""
+
+
+class PoolWorkspace:
+    """What the vendored launches ask their ``workspace`` for, holding nothing between calls.
+
+    :meth:`tile_counter` is the forward's zeroed ``uint32`` ticket counter, new at every call (the kernel leaves it
+    zero). :meth:`scratch` is the backward's ``(dq_accum, lse_log2, dpsum)``: views of one buffer of the workspace
+    pool, key ``("gqa", "bwd_scratch")``, on the current stream.
+    """
+
+    def __init__(self, device: torch.device) -> None:
+        self.device = device
+
+    def tile_counter(self) -> torch.Tensor:
+        return torch.zeros(1, dtype=torch.uint32, device=self.device)
+
+    def scratch(self, total_q: int, num_heads_q: int, num_segments: int) -> tuple[torch.Tensor, ...]:
+        nbytes = cake_jit.scratch_bytes(total_q, num_heads_q, num_segments)
+        buffer = workspace.take(OP, "bwd_scratch", (nbytes,), torch.uint8, self.device)
+        return tuple(
+            buffer[offset : offset + 4 * math.prod(shape)].view(torch.float32).view(shape)
+            for offset, shape in cake_jit.scratch_layout(total_q, num_heads_q, num_segments).values()
+        )
+
+
+def check_seqlens(
+    cu_seqlens_q: torch.Tensor, cu_seqlens_k: torch.Tensor, total_q: int, total_k: int, max_seqlen_k: int
+) -> None:
+    """``ValueError`` unless the call keeps the contract the kernels cannot check.
+
+    Each boundary vector starts at 0, never decreases and ends at its token count, and no segment has more than
+    ``max_seqlen_k`` keys. Reads both vectors on the host: one device synchronization.
+    """
+    split = cu_seqlens_q.numel()
+    bounds = torch.cat((cu_seqlens_q.reshape(-1), cu_seqlens_k.reshape(-1))).tolist()
+    for name, cu, total in (("cu_seqlens_q", bounds[:split], total_q), ("cu_seqlens_k", bounds[split:], total_k)):
+        if not cu or cu[0] != 0 or cu[-1] != total or any(b < a for a, b in zip(cu, cu[1:])):
+            shown = f"{cu[:8]}{' ...' if len(cu) > 8 else ''}"
+            raise ValueError(f"{name} must start at 0, never decrease and end at {total}; got {shown}")
+    keys = [b - a for a, b in zip(bounds[split:], bounds[split + 1 :])]
+    longest = max(keys, default=0)
+    if longest > max_seqlen_k:
+        raise ValueError(
+            f"max_seqlen_k is {max_seqlen_k}, but segment {keys.index(longest)} has {longest} keys: the backward "
+            "would leave key blocks out and give wrong gradients"
+        )
+
+
+class _VarlenGQAAttention(torch.autograd.Function):
+    @staticmethod
+    def forward(  # type: ignore[override]
+        ctx: Any,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        cu_seqlens_q: torch.Tensor,
+        cu_seqlens_k: torch.Tensor,
+        softmax_scale: float,
+        max_seqlen_k: int,
+    ) -> torch.Tensor:
+        out, lse = cake_backend.gqa_attention_fwd(
+            q, k, v, cu_seqlens_q, cu_seqlens_k, softmax_scale=softmax_scale, workspace=PoolWorkspace(q.device)
+        )
+        ctx.save_for_backward(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)
+        ctx.softmax_scale = softmax_scale
+        ctx.max_seqlen_k = max_seqlen_k
+        return out
+
+    @staticmethod
+    def backward(ctx: Any, dout: torch.Tensor) -> tuple[torch.Tensor | None, ...]:  # type: ignore[override]
+        q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k = ctx.saved_tensors
+        dq, dk, dv = cake_backend.gqa_attention_bwd(
+            q,
+            k,
+            v,
+            out,
+            lse,
+            dout.contiguous(),
+            cu_seqlens_q,
+            cu_seqlens_k,
+            softmax_scale=ctx.softmax_scale,
+            max_seqlen_k=ctx.max_seqlen_k,
+            workspace=PoolWorkspace(q.device),
+        )
+        return dq, dk, dv, None, None, None, None
+
+
+def varlen_gqa_attention(
+    q: torch.Tensor,
+    k: torch.Tensor,
+    v: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    *,
+    max_seqlen_k: int,
+    softmax_scale: float | None = None,
+) -> torch.Tensor:
+    """Packed causal GQA attention with autograd: BF16 ``out [T_q, H_q, 256]``; ``q``, ``k`` and ``v`` get BF16
+    gradients.
+
+    ``cu_seqlens_q`` / ``cu_seqlens_k``: int32 ``[B + 1]`` on the device, ``cu[0] = 0``, non-decreasing,
+    ``cu_q[B] = T_q``, ``cu_k[B] = T_k``. ``max_seqlen_k`` (required): a host int at least every segment's key
+    count, as FlashAttention's; a smaller value gives wrong gradients without an error (see the module documentation).
+    ``softmax_scale``: default ``256 ** -0.5``.
+
+    Raises ``ValueError`` for operands the kernels do not take (dtype, head dim, an odd query group, strides, the
+    boundary vectors) and in a process configured for deterministic mode (there is no ``det`` variant),
+    ``TypeError`` when ``max_seqlen_k`` is not an int, ``RuntimeError`` on another GPU. With the option
+    ``cake_gqa_check_seqlens`` it also checks the boundaries and ``max_seqlen_k`` on the host (:func:`check_seqlens`).
+    """
+    config.variant(OP)  # "fast", the only variant; raises in a process configured for deterministic mode
+    if isinstance(max_seqlen_k, bool) or not isinstance(max_seqlen_k, int):
+        raise TypeError(
+            f"max_seqlen_k must be a host int, at least every segment's key count; got {type(max_seqlen_k).__name__}"
+        )
+    if max_seqlen_k < 0:
+        raise ValueError(f"max_seqlen_k must be >= 0, got {max_seqlen_k}")
+    if config.option(OP, "cake_gqa_check_seqlens"):
+        check_seqlens(cu_seqlens_q, cu_seqlens_k, q.shape[0], k.shape[0], max_seqlen_k)
+    scale = HEAD_DIM**-0.5 if softmax_scale is None else float(softmax_scale)
+    return _VarlenGQAAttention.apply(q, k, v, cu_seqlens_q, cu_seqlens_k, scale, max_seqlen_k)

miles_ext_kernels/cake/_runtime/registry.py

@@ -31,6 +31,7 @@
     "rmsnorm",
     "dsa",
     "mhc",
+    "gqa",
 )
 
 # miles_ext_kernels.cake: op modules are its submodules, implementations those of its _impl package

miles_ext_kernels/cake/_vendor/VENDORED.json

@@ -967,6 +967,40 @@
         }
       }
     },
+    "gqa": {
+      "cake_gqa_train": {
+        "cake_revision": "7ab5de0c9d4e42a367ba3e9e218d98208f9bd8fa",
+        "files": {
+          ".gitignore": "e5fd70b13ee2e116ccd3094d1035544ee7f627ed",
+          "README.md": "d150ca19da18df0c2aeccd77668937c41486348e",
+          "__init__.py": "a8171e69c4ff0b4b32f5fa80b0a8aca1b0ff0524",
+          "build.json": "79acfb251bc55693fdad8f8c262d7819069d10e2",
+          "cake_backend.py": "3586d0322a850e7f75333635cd1ff9583c857871",
+          "cake_build.py": "00756137eb159e30a77907eda84fd245123d87ea",
+          "cake_jit.py": "982ebc930de678e9216d7e051ee9796064f89ae8",
+          "csrc/.clang-format": "47a38a93f2dbdd6d944f9dddc3465e3fb65aec29",
+          "csrc/cake_ffi_checks.h": "129288f9432e054dc51b77d93194fb0f2ec20d9c",
+          "csrc/cake_gqa_train_1afdfdde603ff8806b4d_binding.cu": "005e8da7bf8bfc6c20469921b056257da8f1b4e4",
+          "csrc/cake_gqa_train_1afdfdde603ff8806b4d_kernel.cu": "d879f473402e7037a89573c971838e8f6f486252",
+          "csrc/cake_gqa_train_37f4cd45025705f3cfc8_binding.cu": "263c17eb8fdaf6a5db12f53cff290b929c88b19b",
+          "csrc/cake_gqa_train_37f4cd45025705f3cfc8_kernel.cu": "83f14ab04f97d6a4c8f0f38b548935164e8c0c9c",
+          "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_binding.cu": "3dfd95bee83febd5bb32f1aba95df1c549c876f7",
+          "csrc/cake_gqa_train_493c4a28dafa5fc6b2f1_kernel.cu": "11dc829a88a1cd6628c31bdc55b2ac4d00dd1e74",
+          "csrc/cake_gqa_train_6fa13db592acdb66a6c9_binding.cu": "39577c16e0e11e96c9147c2850a0b8717c3180f4",
+          "csrc/cake_gqa_train_6fa13db592acdb66a6c9_kernel.cu": "d6900616e6eaf6118da78a61efd0001d4707ea86",
+          "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_binding.cu": "8cb01566afeb72272e39fcf3e84ff7cc2d8aef90",
+          "csrc/cake_gqa_train_76decc9a5d22dc91ddd1_kernel.cu": "00f1909bf8e29c6087769ba20ecf9477492b5fab",
+          "csrc/cake_gqa_train_a91752e8a2732bf2e246_binding.cu": "7c8ea54ebbd80e7d1305b279296c7027418ea546",
+          "csrc/cake_gqa_train_a91752e8a2732bf2e246_kernel.cu": "574e593a9cabf04738c7a68dc30c8bd06cb379f3"
+        },
+        "source": {
+          "commit": "c5db8d3ccd3fa3842812280f7f208718dc0b8610",
+          "kind": "flashinfer",
+          "path": "flashinfer/experimental/cake_gqa_train",
+          "repo": "https://github.com/flashinfer-ai/flashinfer"
+        }
+      }
+    },
     "kda": {
       "cake_kda_det_train": {
         "cake_revision": "3462767011d1e0de9c84de41bbb90397fa1ebafb",

miles_ext_kernels/cake/_vendor/cake_gqa_train/.gitignore

@@ -0,0 +1,2 @@
+# Prebuilt libraries and their build_info.json sidecar (cake_build.py); built locally, never tracked.
+aot/

miles_ext_kernels/cake/_vendor/cake_gqa_train/README.md

@@ -0,0 +1,91 @@
+# Cake packed-varlen GQA training attention backend (head dim 256, SM100 / SM103 / SM107)
+
+Experimental backend for the forward and first-order backward of causal GQA
+attention over packed variable-length segments, head dimension 256: BF16
+`q [T_q, H_q, 256]`, `k`, `v [T_k, H_kv, 256]` (`H_q / H_kv` even), segments
+packed by `cu_seqlens_q` / `cu_seqlens_k` (int32 on the device, never read on
+the host).  Query row `i` of a segment with `q_len` rows and `k_len` keys sees
+the keys `j <= i + k_len - q_len` (FlashAttention's bottom-right alignment);
+a row that sees no key gets `out = 0`, `lse = -inf` and zero gradients.
+`T_q`, `T_k`, the segment count and every length are run-time values: no
+shape list, no recompilation, no device-to-host read, no device-wide
+synchronization.
+
+Both the API and the backend are experimental: no compatibility guarantee.
+SM100 (B200 / GB200), SM103 (B300 / GB300) and SM107 (R200) are the
+supported architectures; another device raises `RuntimeError`.  The feature does not
+participate in automatic backend selection, autotuning or trace-apply.  Its
+libraries are loaded prebuilt from `aot/` when a builder compiled them and
+JIT-built by FlashInfer from the sources and flags `build.json` records
+otherwise.  The package reads no environment variable.
+
+## Opt-in
+
+```python
+from flashinfer.gqa_train import varlen_gqa_attention, gqa_train_workspace
+
+ws = gqa_train_workspace(q.device)               # optional; reused across calls and shapes
+out = varlen_gqa_attention(
+    q, k, v, cu_seqlens_q, cu_seqlens_k,
+    softmax_scale=0.0625,                        # default 256 ** -0.5
+    max_seqlen_k=max_k,                          # optional host bound of the longest key segment
+    workspace=ws,
+)
+out.backward(dout)                               # BF16 dq, dk, dv
+
+# FP32 dK / dV added into caller-owned FP32 destinations [N_dst, H_kv, 256]
+# through an int32 row map [T_k] (repeated rows accumulate); the callable runs
+# in backward, so nothing is held between the forward and the backward:
+out = varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k,
+                           dkv_acc=lambda: (dk_dst, dv_dst, row_map))
+# FP32 window gradients [T_k, H_kv, 256] stored into caller-owned tensors:
+out = varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k,
+                           dkv_acc=lambda: (dk_win, dv_win))
+```
+
+`varlen_gqa_attention_fwd` / `varlen_gqa_attention_bwd` are the two halves
+without autograd (`dkv_dtype=torch.float32` returns FP32 window gradients;
+`dk_acc`, `dv_acc`, `row_map` add into destinations).  `max_seqlen_k` (a
+host int, never read from the device) must be at least every segment's key
+count, as FlashAttention's `max_seqlen_k`: a smaller value leaves key blocks
+out.  It sizes the backward grid (one CTA pair per 128 keys of the longest
+segment, per key head and segment); without it the bound is `T_k`, which is
+always correct but launches CTAs that retire at once: on a 12-segment layout of
+257K keys the backward call took 8 to 15 % longer without it on SM100, SM103
+and SM107 GPUs.  The backend package
+exposes the same entry points (`varlen_gqa_attention`, `gqa_attention_fwd`,
+`gqa_attention_bwd`, `Workspace`).
+
+## Kernels
+
+| stage | kernel |
+|---|---|
+| `fwd` | O and the natural-log LSE `[H_q, T_q]`: a CTA pair per 256 query rows of one (segment, query head), persistent tile tickets from a self-resetting counter |
+| `bwd_pre` | `dPsum = sum(O * dO)`, `LSE * log2(e)` and the zeroed FP32 dQ accumulator, two query heads per scratch head |
+| `bwd_fp32` / `bwd_rowmap` / `bwd_bf16` | the 5-GEMM backward: a CTA pair per 128 keys of one (segment, key head), looping the key head's query heads two at a time, dK / dV summed in FP32 in tensor memory; FP32 window stores, FP32 row-map adds or BF16 stores |
+| `bwd_post` | `dQ = bf16(acc * softmax_scale)` |
+
+Numerics: FP32 accumulation for the scores, `dP`, dK, dV, the dQ partials and
+the row sums; P and dS are rounded to BF16 only as MMA operands.  dQ and the
+row-map adds use atomic reductions, so they are not bitwise reproducible run
+to run.
+
+## Checkpointing and memory
+
+Autograd saves `(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)` through
+`ctx.save_for_backward` only; `ctx` holds scalars, the `dkv_acc` callable and
+a weak reference to the workspace.  Under
+`torch.utils.checkpoint(..., use_reentrant=False)` a forward retains nothing.
+The backward scratch is FP32 `[H_q / 2, rows_pad, 256]` plus two
+`[H_q / 2, rows_pad]` with `rows_pad = round_up(2 T_q + 128 B, 128)`
+(`gqa_train_workspace_bytes`); a workspace keeps the largest one seen
+(`.nbytes`, `.release()`).  One workspace serves one stream at a time;
+without one, each device stream gets its own default workspace.
+
+## Limitations
+
+- Head dimension 256, BF16 inputs, causal only (bottom-right aligned), no
+  dropout, no sliding window.
+- `H_q / H_kv` must be even.
+- `dout` of the low-level backward must be contiguous (autograd makes it so).
+- `sm_100a`, `sm_103a` and `sm_107a` only.