[Cake 08/11] 08-gqa (review stack)

README.md

@@ -23,6 +23,7 @@ Under `--use-miles-ext-kernels`, miles calls `miles_ext_kernels.add_arguments(pa
 | `kda` (deterministic packed KDA chunked forward + backward: raw forget gate and q/k L2 norm in the kernels, 1-64 heads) | `--kda-backend cake` | none | SM100a / SM103a / SM107a | `tests/kda` (CPU; `test_kda_gpu.py` needs CUDA, its fla parity rows need fla) |
 | `dsa` (DSA sparse-MLA attention forward + backward, lightning-indexer scores; variants det and fast) | `--dsa-attention-backend cake` | `--cake-dsa-variant {det,fast}`, `--cake-dsa-bwd-workspace-bytes`, `--cake-dsa-lazy-build` | SM100a / SM103a / SM107a | `tests/dsa` (CPU: `test_arguments.py`, `test_registry.py`; `test_dispatch.py` needs torch; `test_dsa_gpu.py` needs SM100a / SM103a / SM107a) |
 | `mhc` (Cake four-stream mHC hyper connection: prepare, finish, stream-mean head; det and fast) | `--mhc-backend cake` | `--cake-mhc-variant {det,fast}` | SM100a / SM103a / SM107a (SM107a: hidden 4096 only) | `tests/mhc` (CPU: `test_arguments.py`, `test_vendored_routes.py`; `test_mhc_gpu.py` needs SM100a / SM103a / SM107a) |
+| `gqa` (packed causal GQA attention, head dim 256, forward + backward; dQ is summed with FP32 atomics, so it is not bit-reproducible) | `--full-attention-backend cake` | none (the cake backend is rejected under `--deterministic-mode`) | SM100a / SM103a / SM107a | `tests/gqa` (CPU; `test_gqa_gpu.py` needs CUDA) |
 | `conv1d` (packed causal depthwise conv1d + SiLU, kernel size 4: the GDN / KDA short convs; det + fast) | `--linear-attention-conv-backend cake` | `--cake-conv1d-variant {det,fast}`, `--cake-conv1d-silu-grad-preactivation {bf16,fp32}` | SM100a / SM103a / SM107a | `tests/conv1d` (CPU; `test_conv1d_gpu.py` needs CUDA) |
 | `gated_rmsnorm` (RMSNorm(x) * w * SiLU(z) over rows of 128: the GDN output norm; det + fast) | `--linear-attention-output-norm-backend cake` | `--cake-gated-rmsnorm-variant {det,fast}` | SM100a / SM103a / SM107a | `tests/gated_rmsnorm` (CPU; `test_gated_rmsnorm_gpu.py` needs CUDA) |
 | `rmsnorm` (zero-centred RMSNorm x * rstd * (1 + w): the GDN input layernorm, hidden 256/512/1536/2048/4096/5120/6144; det + fast) | `--linear-attention-input-norm cake` | `--cake-rmsnorm-variant {det,fast}` | SM100a / SM103a / SM107a | `tests/rmsnorm` (CPU; `test_rmsnorm_gpu.py` needs CUDA) |
@@ -32,7 +33,8 @@ faster. Unset, it runs `fast`, or `det` under miles' `--deterministic-mode` (`mi
 `fast` variants of `dsa`, `mhc`, `conv1d`, `gated_rmsnorm` and `rmsnorm` reduce in a
 run-dependent order, so they are not bit-reproducible and `--deterministic-mode` rejects an explicit `fast`. The `gdn`
 `fast` runs in a fixed order: it repeats bit for bit run to run (its bits differ from `det`), so `--deterministic-mode`
-allows it. `kda` ships `det` only. A call that `fast` does not serve runs `det` and logs it once.
+allows it. `kda` ships `det` only; `gqa` ships one variant, not bit-reproducible, which
+`--deterministic-mode` rejects. A call that `fast` does not serve runs `det` and logs it once.
 
 A package applies its parsed arguments in `configure(args)`; miles calls it where it builds a layer on the `cake`
 backend (`miles_ext_kernels.gdn.configure`). miles keeps the values of the arguments registered here in one typed

miles_ext_kernels/__init__.py

@@ -10,6 +10,7 @@
 from miles_ext_kernels.dsa import arguments as dsa_arguments
 from miles_ext_kernels.gated_rmsnorm import arguments as gated_rmsnorm_arguments
 from miles_ext_kernels.gdn import arguments as gdn_arguments
+from miles_ext_kernels.gqa import arguments as gqa_arguments
 from miles_ext_kernels.kda import arguments as kda_arguments
 from miles_ext_kernels.mhc import arguments as mhc_arguments
 from miles_ext_kernels.rmsnorm import arguments as rmsnorm_arguments
@@ -20,6 +21,7 @@
     kda_arguments,
     dsa_arguments,
     mhc_arguments,
+    gqa_arguments,
     conv1d_arguments,
     gated_rmsnorm_arguments,
     rmsnorm_arguments,

miles_ext_kernels/gqa/__init__.py

@@ -0,0 +1,42 @@
+"""
+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 training kernels, head dimension 256, for SM100a / SM103a / SM107a.
+
+The Cake kernels of FlashInfer's ``flashinfer.gqa_train``, vendored under ``_vendor/cake_gqa_train``: forward and
+backward over packed variable-length segments, BF16 ``q [T, H_q, 256]`` and ``k`` / ``v [T, H_kv, 256]`` with
+``H_q / H_kv`` even, causal with FlashAttention's bottom-right alignment per segment. The output and the dK / dV are
+bit-identical run to run; dQ is summed with FP32 atomic adds, so it is not, and deterministic mode rejects the kernels.
+
+miles selects them with ``--full-attention-backend cake`` under ``--use-miles-ext-kernels`` for the full-attention
+layers of the Qwen3.5 / Qwen3-Next models. Importing this package does not import torch: the entry points load on
+first access.
+"""
+
+from importlib import import_module
+
+# public name -> the submodule defining it, imported on first access
+_LAZY = {
+    "SUPPORTED_CAPABILITIES": ".ops",
+    "device_arch": ".ops",
+    "prebuild": ".ops",
+    "varlen_gqa_attention": ".ops",
+}
+
+__all__ = ["SUPPORTED_CAPABILITIES", "device_arch", "prebuild", "varlen_gqa_attention"]
+
+
+def __getattr__(name: str):
+    try:
+        module = _LAZY[name]
+    except KeyError:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
+    value = getattr(import_module(module, __name__), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__():
+    return sorted({*globals(), *__all__})

miles_ext_kernels/gqa/_vendor/VENDORED.json

@@ -0,0 +1,33 @@
+{
+  "fi_commit": "c5db8d3ccd3fa3842812280f7f208718dc0b8610",
+  "fi_repo": "https://github.com/flashinfer-ai/flashinfer",
+  "packages": {
+    "cake_gqa_train": {
+      "cake_revision": "7ab5de0c9d4e42a367ba3e9e218d98208f9bd8fa",
+      "fi_path": "flashinfer/experimental/cake_gqa_train",
+      "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"
+      }
+    }
+  }
+}

miles_ext_kernels/gqa/_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/gqa/_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.

miles_ext_kernels/gqa/_vendor/cake_gqa_train/__init__.py

@@ -0,0 +1,51 @@
+"""
+Copyright (c) 2026 by FlashInfer team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+# Experimental Cake backend: packed-varlen causal GQA attention training with
+# head dimension 256 on SM100 / SM103 / SM107.  Q ``[T_q, H_q, 256]``, K / V
+# ``[T_k, H_kv, 256]`` (BF16, ``H_q / H_kv`` even), segments packed by
+# ``cu_seqlens_q`` / ``cu_seqlens_k`` (int32 on the device), causal with the
+# bottom-right alignment per segment.  The entry point with autograd is
+# ``varlen_gqa_attention``; ``gqa_attention_fwd`` / ``gqa_attention_bwd`` are
+# the two halves.  The registry and the host plan live in ``cake_jit``, the
+# loader, launches and autograd in ``cake_backend``.  Every switch is an
+# explicit keyword argument; the package reads no environment variable.
+
+from .cake_backend import (
+    ARCHES,
+    HEAD_DIM,
+    DkvAccSource,
+    VarlenGQAAttentionFunction,
+    Workspace,
+    default_workspace,
+    gqa_attention_bwd,
+    gqa_attention_fwd,
+    varlen_gqa_attention,
+    workspace_bytes,
+)
+
+__all__ = [
+    "ARCHES",
+    "HEAD_DIM",
+    "DkvAccSource",
+    "VarlenGQAAttentionFunction",
+    "Workspace",
+    "default_workspace",
+    "gqa_attention_bwd",
+    "gqa_attention_fwd",
+    "varlen_gqa_attention",
+    "workspace_bytes",
+]