Repository navigation
[Cake 8/13] gqa: packed causal GQA attention, head dim 256 - #8
Merged
Merged
Conversation
This was referenced Nov 4, 2026
Collaborator
Review of layer 8 (gqa) — suggest approve with nits (document the contract)Confirmed on GPU
Cleared on GPU
Verified on two GB300 nodes (sm_103a), CUDA 13.0, torch 2.14.1+cu130, at the stack head |
Vendor FlashInfer's two Cake short-conv packages at FlashInfer
c5db8d3ccd3fa3842812280f7f208718dc0b8610 into miles_ext_kernels/cake/_vendor/, one level deep:
- cake_causal_conv1d_det_train_backend (flashinfer/cake_causal_conv1d_det_train_backend,
87 files): det, fixed-order weight-gradient reduction, with its context-parallel route
- cake_causal_conv1d_train_backend (flashinfer/experimental/cake/cake_causal_conv1d_train_backend,
75 files): fast, unordered weight-gradient reduction, with its context-parallel route
Both kcp.py modules read one environment variable, a private test hook of the package; the manifest
lists them in allow.env_read. Every file is byte for byte FlashInfer's.
Command (from the repository root, a FlashInfer clone next to it):
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op conv1d \
--package flashinfer/cake_causal_conv1d_det_train_backend \
--package flashinfer/experimental/cake/cake_causal_conv1d_train_backend \
--allow-env-read cake_causal_conv1d_det_train_backend:kcp.py \
--allow-env-read cake_causal_conv1d_train_backend:kcp.py
Co-Authored-By: Claude <noreply@anthropic.com>
Vendor FlashInfer's Cake packed causal GQA training attention package
(head dim 256, SM100a / SM103a / SM107a) from flashinfer-ai/flashinfer
c5db8d3ccd3fa3842812280f7f208718dc0b8610
flashinfer/experimental/cake_gqa_train
byte for byte into miles_ext_kernels/cake/_vendor/cake_gqa_train/, and
record it as the gqa op's block of _vendor/VENDORED.json, with
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op gqa \
--package flashinfer/experimental/cake_gqa_train
run in this repository with a FlashInfer checkout of that commit next to
it. The package has 21 files; their git blobs and its Cake revision
(7ab5de0c9d4e42a367ba3e9e218d98208f9bd8fa) are the ones it had before
this layout, from the same FlashInfer commit. Nothing outside _vendor/
changes; the op that loads the package comes in the next commit.
Co-Authored-By: Claude <noreply@anthropic.com>
The three FlashInfer Cake mHC training packages at c5db8d3ccd3fa3842812280f7f208718dc0b8610,
byte for byte, each one level under miles_ext_kernels/cake/_vendor/ and in the mhc block of
VENDORED.json:
- cake_mhc_det_train_backend (flashinfer/cake_mhc_det_train_backend): det, the forward and a
backward whose parameter gradients are fixed-order reductions; 221 files.
- cake_mhc_train_backend (flashinfer/experimental/cake/cake_mhc_train_backend): fast, copies of
the det forward programs and a backward that adds the dfn partials with FP32 atomics; 181 files.
- cake_mhc_det_train_head_backend (flashinfer/cake_mhc_det_train_head_backend): the model head's
stream mean and its gradient; 12 files.
The package trees are the ones this op shipped before the layout change (same git trees). No
vendored module reads the environment or imports flashinfer, so the entries list no allow hooks.
Vendored by running, in this repository with a FlashInfer checkout of that commit next to it:
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op mhc \
--package flashinfer/cake_mhc_det_train_backend \
--package flashinfer/cake_mhc_det_train_head_backend \
--package flashinfer/experimental/cake/cake_mhc_train_backend
Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.cake.mhc declares the op (torch-free): variants det and fast, selected by
--mhc-backend cake, --cake-mhc-variant, the three vendored packages, and the entry points
mhc_prepare, mhc_finish, mhc_head, is_supported, supported_hidden_sizes, HIDDEN_SIZES and STREAMS,
loaded from miles_ext_kernels.cake._impl.mhc on first use. Its validate rejects a --hidden-size
the kernels do not run.
miles_ext_kernels.cake._impl.mhc.ops routes a connection to the det or the fast package by the
variant configured when mhc_prepare runs; the handle carries it to mhc_finish and the backward.
Changes in behaviour:
- A process that never configured mhc runs det, and configure(deterministic_mode=True) without a
variant runs det (it ran fast). An explicit fast in deterministic mode is rejected.
- The variant comes from the configuration only: torch.use_deterministic_algorithms(True) no longer
turns a fast call into det; the fast backward then raises torch's alert for an operation without
a deterministic implementation. The per-call variant argument is gone.
- prebuild("mhc", *, device=None, max_workers=None, hidden, sinkhorn_iters) takes both model
parameters, without defaults, and builds every library of both variants and the head at that
shape. A Sinkhorn count other than 20 runs run-time-count programs of the coefficient and
token-owner passes, which a prebuild at the default count left to compile inside the first call.
The epsilons are launch arguments and select no library.
- BF16 fn, scale and base (FSDP2 with param_dtype=torch.bfloat16) are cast to FP32 at the call,
which is exact; their gradients return in BF16. They raised ValueError.
- Workspaces come from the runtime's workspace pool: each launch of the forward, the statistics
replay and the backward takes a buffer when it runs (roles det-fwd, det-bwd, fast-fwd, fast-bwd,
sized for the call's power-of-two token-count bucket), so they stay under
--cake-workspace-pool-bytes and miles_ext_kernels.release() frees them. The packages' own
per-device default workspaces, which lived for the whole process, are no longer allocated.
- release() drops the launch plans and the loaded libraries; the next call reloads the libraries
from the cache.
The packages resolve every launch workspace through one function, which the op replaces with the
pooled one. The kernels' work tickets, the first 256 bytes of every workspace layout, must be zero
before a launch, and the content of a workspace is undefined: every launch zeroes them first (256
bytes on the stream of the call), from the pool or not.
Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.cake.gqa declares the op (torch-free) and
miles_ext_kernels.cake._impl.gqa runs FlashInfer's cake_gqa_train
package through the runtime's loader: varlen_gqa_attention, packed
causal GQA attention over variable-length segments, BF16 q
[T_q, H_q, 256] and k / v [T_k, H_kv, 256] with an even number of query
heads per key/value head, forward and backward, on SM100a, SM103a and
SM107a. miles selects it with --full-attention-backend cake.
Variants. The op ships one variant, fast: the output, the LSE and dK /
dV repeat bit for bit run to run, dQ is summed with FP32 atomic adds.
There is no det variant, so the runtime rejects the op under
--deterministic-mode, and a process configured with
deterministic_mode=True raises at the call. A process that never
configured the op runs fast.
max_seqlen_k is now a required keyword-only host int and a documented
hard contract: at least every segment's key count, as FlashAttention's.
It sizes the backward grid, so a smaller value leaves key blocks out of
the backward and gives wrong dQ / dK / dV without an error, while the
output is unaffected. The kernels never read the boundaries on the
host; the new option --cake-gqa-check-seqlens (configure
cake_gqa_check_seqlens=True) checks at every call that cu_seqlens_q /
cu_seqlens_k start at 0, never decrease and end at the token counts,
and that no segment has more keys than max_seqlen_k, at the price of
one device synchronization per call. It is off by default.
The implicit None bound is gone: it sized the backward grid by T_k,
which is correct but launches CTA pairs that only retire, and made the
backward slower the more segments a call packs. A caller without the
longest segment can still pass T_k explicitly. The output, the LSE and
dK / dV do not depend on the bound given.
check_supported(num_heads_q, num_heads_kv, head_dim=256) and
is_supported(...) tell, without torch, whether the kernels run a head
layout. validate_args applies them to the heads one tensor-parallel
rank of miles holds (num_attention_heads / TP query heads over
max(num_query_groups, TP) / TP key/value heads, head dim from
--kv-channels or hidden_size / num_attention_heads), so a layout such as
24 query heads over 4 groups at TP 8 (3 query heads over 1 key/value
head per rank) is rejected at startup rather than at the first forward.
Memory. Each backward takes its scratch from the runtime's workspace
pool on the current stream, and each forward zeroes a fresh tile
counter; nothing goes through the package's per-stream default
workspaces, which grow to the largest call and live as long as the
process. release() therefore frees all of it, and calls on two streams
never share a buffer. prebuild("gqa") builds the four libraries of the
forward and the BF16 dK / dV backward for the device's architecture; it
takes no model parameter.
The op's names move from miles_ext_kernels.gqa:
SUPPORTED_CAPABILITIES, device_arch and prebuild are the runtime's
ARCH_BY_CAPABILITY, device_arch and prebuild("gqa"). The README lists
the op in its Ops and Prebuild tables.
Co-Authored-By: Claude <noreply@anthropic.com>
tests/cake/gqa/, CPU: - test_arguments.py (torch-free): the op registers one variant and no variant flag; --cake-gqa-check-seqlens is off by default and needs --full-attention-backend cake; --deterministic-mode (also through the typed configuration) rejects the op; check_supported / is_supported on head layouts; validate_args on the heads of one tensor-parallel rank (24 query heads over 4 groups at TP 8 and 8 over 8 at TP 8 rejected, 16 over 4 at TP 8 accepted, head dim from --kv-channels or hidden_size / num_attention_heads); configure and option. - test_gqa_cpu_plan.py (torch-free): every stage the entry point launches has one library per architecture, which the oldest toolkit the runtime accepts for it (MIN_CUDA) builds; the grids cover every tile; the backward grid covers the key blocks of max_seqlen_k and no more; the scratch layout. - test_gqa_host.py (torch on the CPU): deterministic mode raises at the call and configure() lets the next call through; max_seqlen_k is a required host int; check_seqlens accepts the contract and rejects a bound below the longest segment and malformed boundaries; the option runs the check before any launch; with the pool on, the backward scratch is one pooled buffer in the package's layout, and at the default bound, 0, each backward takes a buffer of its own that the pool does not keep; prebuild takes no model parameter, raises NotImplementedError on another GPU and RuntimeError when the toolkit is older than CUDA 13.4 on SM107a. GPU (skip without SM100a / SM103a / SM107a): - test_gqa_gpu.py: FlashInfer's suite layouts and miles' head pairs against an FP64 reference (assert_rel_rms) within FlashInfer's budget capped at twice the error measured on a B200, noted next to it: FlashInfer's budget for out, dQ, dK and dV (1.08 to 1.12 times the measured error), twice the measured error for the LSE, whose errors are at FP32 rounding; out, LSE, dK and dV bit for bit run to run in one process and against a new process, for any valid max_seqlen_k (exact, larger, T_k), on alternating default and side streams, and pooled against the pool turned off; dQ within its FP32 reordering; the check rejects max_seqlen_k 1024, 1920 and 2047 on 2048-key segments and changes no bit when valid; document isolation; checkpointing; no device synchronization in steady state (and one with the check); with the pool on, the pooled scratch and release(); rejected operands; two GPUs give the same bits (multi_gpu). - test_prebuild_gpu.py: ROUTES for every call-time path (exact and T_k max_seqlen_k, groups 2 / 4 / 6 / 8, keys fewer than queries, empty segments, long keys, one segment, no query, no key, a softmax scale, the check option, a side stream, the pool on) after one prebuild under no_build(); prebuild returns the four libraries; a fresh process loads only those and never imports flashinfer. Co-Authored-By: Claude <noreply@anthropic.com>
CPU (no GPU): the arguments and the variant rules through the hooks miles calls, deterministic mode without a variant included; the vendored packages against the routes, constants and hidden sizes the op relies on; what prebuild builds per Sinkhorn count, checked against the programs the host rules select at every token count; BF16 parameters reaching the kernels as FP32; the variant read at each call; with the pool on, every launch workspace taken from the pool, its tickets zeroed at every launch, and at the default bound, 0, none kept; what release drops. GPU: det and fast against an FP64 reference at every hidden size the device serves, the fast outputs against det's bits, another Sinkhorn count and other epsilons, BF16 parameters, det bit for bit in a new process, on a side stream and on a second GPU (multi_gpu), with the pool on: the bits under any pool bound, from a buffer a larger call sized and with interleaved connections, release; the head, and test_prebuild_gpu.py: ROUTES over every call-time switch at the Sinkhorn count 7, checked from an empty library cache. Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels/cake/conv1d.py declares the op: variants det and fast, selected by
--linear-attention-conv-backend cake, the option --cake-conv1d-silu-grad-preactivation, the two
vendored packages and the entry points causal_conv1d, causal_conv1d_kcp, causal_conv1d_kcp_multi and
check_supported. miles_ext_kernels/cake/_impl/conv1d/ implements them on the runtime's loader, config
and prebuild. The op is appended to OPS and to the README's Ops and Prebuild tables.
Changes in behavior:
- A process that never configured the op runs det; it ran fast. miles' command line still defaults
to fast: validate_args resolves --cake-conv1d-variant to fast and writes it back for configure.
- Each call reads its variant and SiLU' policy when it runs; the backward runs those of its forward.
- The entry points take no variant argument: only configure selects the variant, so a process
configured with deterministic_mode runs det on every call. A call could pass variant="fast" before.
- --cake-conv1d-silu-grad-preactivation defaults to bf16; it had no default. Without
--linear-attention-conv-backend cake, an explicit bf16 is now accepted (it raised before); any
other value still requires the cake backend.
- causal_conv1d without cu_seqlens builds the prefix sums of its one document on the device. It copied
them from pageable host memory before, which synchronized the host on every call.
- check_supported(..., has_bias=True) raises: the entry points take no bias, so a biased layer
would lose its bias.
- prebuild("conv1d", weight_dtype=...) builds every library of both variants, both SiLU' policies and
the context-parallel route. It replaces prebuild(dtype) and prebuild_kcp().
- release() drops the packages' host plans.
Co-Authored-By: Claude <noreply@anthropic.com>
Move the conv1d tests to tests/cake/conv1d/ and run them through the runtime (configure, variant, prebuild, release). - CPU: the arguments and the configuration (det unconfigured, fast from miles' command line, an explicit default SiLU' policy without the cake backend), the variant and SiLU' policy read at each call, no entry point taking a variant (det on every call under deterministic_mode), the one-document cu_seqlens built on the device, the rejected bias, prebuild's arguments and its libraries against the packages' program selection on sm_100a, sm_103a and sm_107a, release, the vendored packages and their loading without flashinfer. - GPU: the accuracy and determinism tests as before, each variant selected with configure, and new ones: det bitwise in a new process, cu_seqlens=None without a host synchronization or a wait for queued work, calls on a side stream interleaved with the default stream, an unconfigured process running det, and test_prebuild_gpu.py, whose ROUTES cover every call-time path and run under no_build() after one prebuild, for each taps dtype in a new process with an empty library cache of its own. - The multi-GPU context-parallel test is marked multi_gpu. Co-Authored-By: Claude <noreply@anthropic.com>
The four Cake KDA training packages of FlashInfer 7a13d7ba9b75c9bb9c02452df36e3f34417b5d57, byte for byte, one
level deep under miles_ext_kernels/cake/_vendor/ and recorded under ops.kda of VENDORED.json:
- cake_kda_det_train (flashinfer/experimental/cake_kda_det_train): det, single rank.
- cake_kda_det_train_kcp (flashinfer/experimental/cake_kda_det_train_kcp): det, context parallel.
- cake_kda_train (flashinfer/experimental/cake/cake_kda_train): fast (deterministic=False), single rank.
- cake_kda_train_kcp (flashinfer/experimental/cake_kda_train_kcp): fast, context parallel.
One FlashInfer commit for the family, as the context-parallel packages import their single-rank siblings
(..cake_kda_det_train and ..cake.cake_kda_train; the runtime resolves the second to _vendor.cake_kda_train).
7a13d7ba rather than c5db8d3c: it carries FlashInfer #8629, a correctness fix of the det single-rank package (with
the in-kernel q / k L2 norm, its replay backward over a batched arena, KdaDetArena(batch_tokens=...), read up to 63
rows of the next document group before they were written), and the other changes of these packages since c5db8d3c:
#8552, #8620, #8625, #8630 and #8652 (performance; #8625 runs the fused L2 norm's GLM intra stage with another
program on SM103a, #8620 sums the fast packages' raw-gate parameter gradients in FP64 on SM100a / SM103a). The entry
points of this repository pass no arena, so #8629 does not change their results. The later commits up to FlashInfer
main change none of the four directories.
cake_kda_train/cake_backend.py keeps two FlashInfer hooks, listed under allow: it reads CAKE_KDA_CKPT_HANDOFF
(the checkpoint hand-off override of its chunk_kda), and _jit_nvcc_version imports flashinfer.jit.cpp_ext inside
its body to name the toolkit of the libraries; the kda implementation rebinds that function to the runtime loader's
nvcc when it is imported.
Vendored with, from the root of this checkout next to a FlashInfer clone at ../flashinfer:
python tools/vendor_cake.py fi --checkout ../flashinfer --commit 7a13d7ba9b75c9bb9c02452df36e3f34417b5d57 \
--op kda \
--package flashinfer/experimental/cake_kda_det_train \
--package flashinfer/experimental/cake_kda_det_train_kcp \
--package flashinfer/experimental/cake/cake_kda_train \
--package flashinfer/experimental/cake_kda_train_kcp \
--allow-env-read cake_kda_train:cake_backend.py \
--allow-flashinfer-import cake_kda_train:cake_backend.py:_jit_nvcc_version
Co-Authored-By: Claude <noreply@anthropic.com>
FlashInfer c5db8d3ccd3fa3842812280f7f208718dc0b8610: the SiLU-gated RMSNorm training packages, deterministic
(cake_rmsnorm_silu_gate_det_train_backend, 14 files) and non-deterministic (cake_rmsnorm_silu_gate_train_backend,
30 files), each one level deep under miles_ext_kernels/cake/_vendor/, and their block ops.gated_rmsnorm in
_vendor/VENDORED.json. Written from a FlashInfer checkout by:
python tools/vendor_cake.py fi --checkout ../flashinfer --commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 \
--op gated_rmsnorm --package flashinfer/cake_rmsnorm_silu_gate_det_train_backend \
--package flashinfer/cake_rmsnorm_silu_gate_train_backend
Co-Authored-By: Claude <noreply@anthropic.com>
Vendor FlashInfer's Cake RMSNorm training packages of flashinfer-ai/flashinfer eb3eb68d7deebbe619e05d93ecfd9796e4eaf2e3 into miles_ext_kernels/cake/_vendor/, byte for byte, as the packages of the op rmsnorm in _vendor/VENDORED.json: - cake_rmsnorm_det_train_backend (flashinfer/cake_rmsnorm_det_train_backend, 582 files): the deterministic forward and backward. - cake_rmsnorm_train_backend (flashinfer/cake_rmsnorm_train_backend, 581 files): the backward with a non-deterministic weight gradient in some token-count windows; it runs its deterministic sibling for deterministic=True, so the two packages move together. eb3eb68d includes flashinfer-ai/flashinfer#8512: on SM100a and SM103a the deterministic backward of the smallest token counts sums each weight-gradient column in a fixed row order, so its bits no longer depend on the SM count. The files and their blobs are the ones the previous rmsnorm package vendored from the same commit. python tools/vendor_cake.py fi --checkout ../flashinfer --commit eb3eb68d7deebbe619e05d93ecfd9796e4eaf2e3 \ --op rmsnorm --package flashinfer/cake_rmsnorm_det_train_backend \ --package flashinfer/cake_rmsnorm_train_backend Co-Authored-By: Claude <noreply@anthropic.com>
Vendor the DSA kernel packages into miles_ext_kernels/cake/_vendor/ with tools/vendor_cake.py, byte for byte, and record them under ops.dsa of VENDORED.json: - cake_dsa_train: FlashInfer's DSA training kernels for 64 local query heads (the fast variant: dKV and dk_rope accumulate in FP32 with red.global). - cake_dsa_indexer_det: FlashInfer's lightning-indexer scores and deterministic exact top-k. - cake_dsa_sparse_det: the DSA det version, the generated sparse attention forward / backward and lightning-indexer logits / backward kernels of radixark/miles#3599 (86 sources and registry.json; no floating-point atomics, key-side gradients reduced from FP32 partials in a fixed 4-way tree). The tool writes the TVM-FFI host surface in place of the upstream torch-extension one and a build.json generated from the registry, so these kernels build with nvcc and load with tvm-ffi through the same loader as the FlashInfer packages; SM107a compiles the SM103a records. Both FlashInfer packages come from c5db8d3, the commit these kernels were validated at: cake_dsa_train is byte for byte the same at 7a13d7ba, and cake_dsa_indexer_det differs there only by a later performance change (flashinfer-ai/flashinfer#8691). Commands (run from the repository root, with a FlashInfer checkout at ../flashinfer and a miles checkout at ../miles): python tools/vendor_cake.py fi --checkout ../flashinfer \ --commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op dsa \ --package flashinfer/experimental/cake_dsa_train \ --package flashinfer/experimental/cake_dsa_indexer_det python tools/vendor_cake.py miles-pr --checkout ../miles \ --pr 3599 --commit a1b763c82360afeb00c551aeb7ea93b30870e5bf \ --op dsa --name cake_dsa_sparse_det \ --path miles/kernels/attention/dsa/cake \ --include csrc --include registry.json \ --exclude csrc/cake_host_shim.h --arch-from sm_107a=sm_103a Co-Authored-By: Claude <noreply@anthropic.com>
The KDA op of the cake namespace: miles_ext_kernels/cake/kda.py declares it (OpSpec, selected by --kda-backend
cake), miles_ext_kernels/cake/_impl/kda/ implements it on the four vendored packages, OPS gains "kda" and the
README its Ops and Prebuild rows. The entry points are those of the kda family before this layout: chunk_kda,
chunk_kda_fwd, chunk_kda_bwd, and partition_kcp, build_context_kcp, chunk_kda_kcp, chunk_kda_fwd_kcp,
chunk_kda_bwd_kcp, reduce_gate_param_grads_kcp. Each call reads config.variant("kda") when it runs: det unless
configured.
Variants. det (cake_kda_det_train, cake_kda_det_train_kcp) repeats bit for bit. fast (cake_kda_train,
cake_kda_train_kcp, deterministic=False) runs the S1 gate-gradient programs on some backward rows of the activated
gate mode and, in the raw gate mode on SM100a / SM103a, reduces the gate parameters' gradients in FP64 in no fixed
order. fast has no in-kernel q / k L2 norm: such a call runs det and logs it once (config.warn_runs_det). Every
miles call takes that L2 norm, so cli_variants is ("det",): --cake-kda-variant is registered hidden, validate_args
writes det back for every process, and a fast choice from miles, which would change no program, is refused by
argparse. Direct callers configure either variant. A backward reads its saved set's package (a saved set is read by
the package that wrote it); saved=None replays the forward in the package of the call. Every call hands its package
use_qk_l2norm_in_kernel as given, so the backward of a fast saved set asked for the L2 norm, a mode its forward did
not run, is the fast package's ValueError rather than a gradient without the norm. Partitions and contexts hold no
variant.
Operands as miles hands them over. beta may be FP32 (miles gives flash-linear-attention the FP32 sigmoid) and v a
column slice of the packed projection: chunk_kda* round beta to BF16 and make v dense before the packages' checks,
inside the autograd graph, and return dbeta in beta's dtype. Before, both raised ValueError in the packages.
Prebuild. prebuild(num_heads=...) checks 1 <= num_heads <= 64 (an int, not a bool) before anything else and then
builds every library of the four packages for the device's architecture (188 / 186 / 189 on SM100a / SM103a /
SM107a): the record a launch takes depends on its token and document counts, its modes, its path and, under context
parallelism, the rank's shard, so no smaller set holds for every call of a layer. This replaces the family's
route_libraries probe, which planned a fixed list of problems and dropped every problem it could not plan (an
unknown gate mode gave an empty list, not an error). release() drops the packages' per-stream entry pools and the
captured graphs of a loaded context-parallel package.
The fast package's _jit_nvcc_version hook (its checkpoint hand-off keys row records on the toolkit) is rebound to
the runtime loader's nvcc when the implementation is imported: flashinfer is never imported. --kda-backend cake
rejects context parallelism (--context-parallel-size > 1) and a layout other than thd at startup, as before.
Co-Authored-By: Claude <noreply@anthropic.com>
The output norm of miles' GDN layers, RMSNorm(x) * weight * SiLU(z) over rows of 128, becomes the op
miles_ext_kernels.cake.gated_rmsnorm: OP declares its two variants, the miles selector
--linear-attention-output-norm-backend cake, its two vendored packages and the entry points rmsnorm_silu_gate and
check_supported; miles_ext_kernels.cake._impl.gated_rmsnorm implements them.
- A call reads the configured variant when its forward runs, and its backward runs the same package. A process that
never configured the op runs det (it ran fast); validate_args still resolves miles' command-line default to fast,
and deterministic mode resolves det and rejects an explicit fast. While torch.use_deterministic_algorithms(True)
is in effect, which the non-deterministic backward refuses, fast runs det and logs it once.
- prebuild(device=None, max_workers=None), reached through cake.prebuild("gated_rmsnorm"), builds every library of
both packages on the device's architecture, every row-count route of each, so no call compiles afterwards.
release() has nothing to drop.
- rmsnorm_silu_gate documents when a gate that is a column slice of a wider projection is read in place, and that
the weight gradient sums over the rows of one call: under tensor parallelism the caller all-reduces it over the
tensor-parallel group.
- registry.OPS lists gated_rmsnorm; the README's Ops and Prebuild tables list its rows.
Gone with the old package: the module-level VARIANT, the op's own arguments module and configure, the per-call
variant keyword, and SUPPORTED_CAPABILITIES (miles_ext_kernels.cake.ARCH_BY_CAPABILITY names the same devices).
Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.cake.rmsnorm declares the op: selected by --linear-attention-input-norm cake, variants det and fast
(--cake-rmsnorm-variant), the two vendored packages. miles_ext_kernels.cake._impl.rmsnorm implements it:
ZeroCenteredRMSNorm, zero_centered_rmsnorm, check_supported, supported_hidden_sizes, prebuild(hidden_size=...) and
release. The op is appended to OPS and to the README's Ops and Prebuild tables.
Changes from the previous rmsnorm package:
- A call reads the configured variant when it runs (det in a process that never configured the op), and its backward
runs the variant of its forward. The module no longer reads or pins a variant at construction; its variant=
argument is gone, configure is the one switch.
- Constructing the module compiles nothing. prebuild("rmsnorm", hidden_size=H) builds every library of both variants
for H (forward and backward, every token-count window), so a variant configured after a module was built runs at
its next call without a build.
- The weight is allocated where torch's factory functions allocate (a with torch.device("meta") block,
torch.set_default_device), else on the current CUDA device; device="meta" builds without touching CUDA. The device is
checked when it is a CUDA device; a hidden size the kernels do not serve is refused on any device.
- reset_parameters() zeroes the weight, also after to_empty().
- sequence_parallel= (default None) is set on the weight when given, at construction and by reset_parameters, as
Transformer Engine's norms set it, so Megatron all-reduces its gradient across the tensor-parallel ranks under
sequence parallelism. None leaves the attribute to the caller: reset_parameters keeps a mark set on the weight.
- An unsupported device raises NotImplementedError (was ValueError), as the runtime's device_arch does.
- The package's configure, prebuild(hidden, device, variant=) and SUPPORTED_CAPABILITIES are replaced by
miles_ext_kernels.configure, miles_ext_kernels.cake.prebuild("rmsnorm", hidden_size=...) and
miles_ext_kernels.cake.ARCH_BY_CAPABILITY.
fast still runs det, logged once, under torch.use_deterministic_algorithms(True), which the non-deterministic package
refuses.
On SM107a, hidden size 4096 launches two programs above the 232448 B shared-memory opt-in (271360 B, CUDA 13.4's
oversized shared-memory mode): det's backward reduce above 4096 tokens and its copy in fast's backward from 4097 to
24064 tokens. A process whose CUDA runtime is older than 13.4, such as the 13.0 runtime wheel a cu130 torch installs,
rejects that mode at the launch ("cudaFuncSetAttribute failed: invalid argument"), and the error it leaves pending
fails the next CUDA call too. The op now checks the loaded runtime (libcudart of the toolkit's major) and the
driver's oversized ceiling for every program of the width that sets the mode, at construction on a CUDA device,
check_supported, prebuild and the first call of a width on a device, and raises RuntimeError naming the programs,
the runtime and the remedy before any launch. Other widths and SM100a / SM103a launch no such program. The README's
Ops and Prebuild rows state the limit.
Co-Authored-By: Claude <noreply@anthropic.com>
Add the DSA op: miles_ext_kernels/cake/dsa.py (its OpSpec, variants det and fast, --cake-dsa-bwd-workspace-bytes and --cake-dsa-lazy-build) and its implementation under miles_ext_kernels/cake/_impl/dsa/: - ops.py: sparse_attention, the lightning indexer, the det forward and backward entry points, prebuild and release. The variant is read when a call runs; a call the fast variant does not serve (other than 64 heads, a sink, another d_qk) runs det and logs it once per reason. - native.py: the DSA det version (cake_dsa_sparse_det) built and loaded through the runtime loader, one library per stage and architecture. Every argument binds by its exported name; a TMA source binds as a tensor whose descriptor the stage's launcher encodes on the host. - kernels.py, reduction.py, layouts.py, indexer_topk.py: the det host side, the fixed-order reduction, the layout adapters and the det indexer top-k (cake_dsa_indexer_det). Register the op in OPS and add its rows to the README's Ops and Prebuild tables. Behaviour, compared with the previous layout: - sm_scale is required: the d_qk ** -0.5 default is not the scale of the absorbed MLA these kernels serve. - A call takes no variant argument; configure() selects it. - A batch-local bshd index outside its batch element's keys is an empty slot; it no longer reads the next element's keys. So is an indexer index past the keys. - --cake-dsa-bwd-workspace-bytes bounds the det backward's partials and the index scratch of their sort (one sort of every chunk when it fits, else one per chunk); the chunking, and with it the det bits, follows from it. - The fast backward receives the packed segment count (num_segments; the batch for bshd) to plan its key passes; the bits do not change. - A call without query rows returns empty outputs and zero gradients, forward and backward, attention and indexer (and the indexer's top-k selection, with or without key rows). - The indexer serves 8, 16, 32 or 64 heads of dimension 128 in its logits, its backward and the selected scores alike; 24 is rejected before any launch. - A selected key with a -inf score (outside its query's window) gets no gradient, even a NaN one: the backward sees it as an empty slot. Co-Authored-By: Claude <noreply@anthropic.com>
tests/cake/kda/: the kda family's tests on the cake layout, with per-tensor accuracy budgets. - test_arguments.py (torch-free): the spec, the hidden --cake-kda-variant (det written back, fast refused), the settings validate_args rejects, configure for direct callers, an import without torch. - test_kda_variant.py (CPU, torch): the package of every entry point under det and fast, the in-kernel L2 norm call (logged once), the backward of a saved set of either package (a fast set asked for the in-kernel L2 norm is the fast package's ValueError), the beta / v seam and its gradients, the head-count check before any package or build runs, the fast toolkit hook, that fast plans programs det does not ship (S1 in the activated gate mode on every architecture, R1c in the raw gate mode on SM100a / SM103a), that prebuild builds the same 188 / 186 / 189 libraries for every head count, and release(). - test_kda_gpu.py: det in miles' call (q / k / v column slices, raw gate, in-kernel L2 norm) against the exact FP64 recurrence, every output and gradient within twice the largest error measured (_budgets.py, one B200 and one GB300): the token tensors in relative RMS error (a 1.02 scale of out, dq, dk, dv or dbeta fails it), dA_log and ddt_bias in error over the root sum of squares of the terms of their exact sums, free of the shape's cancellation; bit-identical repeats, packed views and an FP32 beta against contiguous BF16 operands, autograd / saved set / replay / checkpoint bits, default and side stream bits, a fresh process's bits, and flash-linear-attention parity on every shape (each tensor within twice fla's error on the same inputs). - test_kda_fast_gpu.py: fast against the same reference within twice det's error on the same inputs (no bit checks: fast promises none), its forward / backward split and checkpointed node, saved sets of either package, and the in-kernel L2 norm call that runs det giving det's bits. - test_kda_kcp_gpu.py: the context-parallel entry points on 2 and 4 GPUs, marked multi_gpu, so a one-GPU run deselects it (-m "not multi_gpu") instead of reporting 20 skips; with fewer GPUs than a group needs it skips with that count. - test_prebuild_gpu.py: ROUTES over every call-time switch (gate and L2 norm modes, packed views and FP32 beta, explicit forward and backward, replay, both checkpoints, the segmented route with saved states, 1 to 64 heads, more than 256 documents, a token sweep with empty documents, the context-parallel entry points on a one-rank partition with host cu_seqlens), checked after one prebuild with no build allowed. The elementwise gate the family used (atol = rtol = 1e-2 on every tensor) let a 2 % error of out, dv or dbeta pass; the relative RMS budgets do not, and the elementwise gate is not carried over. dg is exempt from the 1.02 detection check, as its own error is about 1 %. dA_log and ddt_bias are sums over every token of a head. Their relative error is set by how far the terms cancel, which varies with the shape (on one B200 and one GB300, 0.8 % to 29 % for dA_log and 4.1 % to 6.3 % for ddt_bias), and dA_log has one value per head, so a relative budget per shape is loose at one head. Over the root sum of squares of the terms (_reference.gate_param_terms) the error is 0.6 % to 3.3 % for dA_log and 4.7 % to 6.3 % for ddt_bias, its largest at one head; fla's on the same inputs is 1.2 % to 2.8 % and 4.6 % to 5.6 % (the fla parity rows print both). These budgets are twice the largest. Co-Authored-By: Claude <noreply@anthropic.com>
tests/cake/gated_rmsnorm/: - test_arguments.py (CPU): the declaration and its variant flag; an unconfigured process and an unconfigured call run det; miles' resolution (fast by default, det in deterministic mode, explicit values kept, fast rejected in deterministic mode, the flag rejected without the backend); the variant read at every call and det under torch's deterministic algorithms, with the packages replaced by recorders; check_supported's rejections before any device query; import and configuration without torch. - inputs.py: miles' GDN layout. The gate is a column slice of the fused [q | k | v | z | b | a] projection, whose width (2 * Hk * 128 + 2 * Hv * 128 + 2 * Hv, padded to a multiple of 8) is the gate's token stride: the kernels read the gate in place when it is a multiple of 16 elements and copy it otherwise. Every forward + backward checks the documented result types: y BF16 of x's shape, every gradient BF16, dw of shape [128]. - test_gated_rmsnorm_gpu.py: every value-head count (12, 16, 32, 48) with the gate read in place and copied, the path asserted, among them 16 key / 32 value heads (width 12352, in place) and its twin with 8 more columns (12360, copied), at 1 to 32771 tokens: y, dx and dz within one BF16 ulp of FP64 and dw within det's bound, every gate failing a 2 % error; det bitwise three times in one process and once in a new process; a contiguous gate the same bits; fast's y, dx and dz det's bits. Also calls alternating between the default and a side stream, checkpointing, accumulation, a tensor-parallel head split (each rank's dw covers its own heads, their sum is the layer's) and the toolkit check. - test_prebuild_gpu.py: ROUTES (both gate paths at every head count, a 2-D input, a token sweep that reaches every row-count module of the non-deterministic package, checked against the device's SM count, a weight without a gradient, torch's deterministic algorithms) under check_prebuild; prebuild's libraries and idempotence; configure after a call switches the next call without a build; no flashinfer module is imported. Co-Authored-By: Claude <noreply@anthropic.com>
tests/cake/rmsnorm/: - test_arguments.py (torch-free): registration, variant resolution with and without miles' deterministic mode, rejections, configure, det in an unconfigured process. - test_module.py (CPU and meta tensors, the kernels replaced by a recorder): construction on the meta device, reset_parameters after to_empty, the sequence_parallel attribute (none by default; set from the argument; a mark the caller set kept by reset_parameters) and Megatron's selection rule, refusal of widths the kernels do not serve, construction without a build, the variant each call reads, fast running det under torch's deterministic algorithms, the arguments of prebuild (its refusal without a CUDA device in a child process that sees none, so it runs on GPU nodes too); the programs that launch in the oversized shared-memory mode (SM107a at 4096 only) and the refusal of their width, with the CUDA runtime and driver facts replaced (runtime 13.0, none loaded, no oversized ceiling, a ceiling too small). - test_rmsnorm_gpu.py: accuracy against FP64 at FlashInfer's budgets; det bit for bit three times in one process, once in a fresh process, at the smallest token counts and on a side stream (hidden sizes 4096 and 5120); fast's y and dx against det; strided rows, checkpointing, no tokens, a module built on the meta device, the toolkit check, a backward that runs its forward's variant. A width this process cannot launch (SM107a at 4096 without a CUDA 13.4 runtime) is refused instead of run (_refusal.py); the refusal holds exactly where a fresh process with the check disabled fails to launch the width, and with the runtime taken as 13.0 SM107a refuses 4096 at check_supported, prebuild, construction and the first call, while every other width and architecture runs. - test_prebuild_gpu.py: ROUTES (every token-count window of both packages, [T, H] input, strided rows, FP32 weight, no tokens, checkpointing, torch's deterministic algorithms) for every width after one prebuild, with no build allowed (prebuild of a width this process cannot launch raises its refusal); a variant configured after construction runs without a build; no flashinfer module is imported. - test_sequence_parallel_gpu.py (multi_gpu, 2 GPUs): under tensor and sequence parallelism the all-reduce rule selects the weight, and the ranks keep one bit-identical weight that follows the single-GPU run. Co-Authored-By: Claude <noreply@anthropic.com>
Add tests/cake/dsa/: - test_arguments.py, test_package.py, test_dispatch.py and test_indexer_topk.py (CPU): the arguments and configuration, the three packages and the CUDA 13.4 floor of SM107a whatever a package records, the stage selection, the argument binding of every det stage (TMA sources included), the calls fast serves, the required sm_scale, the segment count, the index ranges, the det backward's memory plan, calls without query rows (the indexer forward and backward and its top-k selection included), the indexer head counts and the gradient mask of the selected scores. - test_dsa_gpu.py and test_indexer_topk_gpu.py: accuracy against FP32 and FP64 references at every head count from 4 to 128, det bits repeated in one process, in a fresh process, on a side stream and on a second GPU, the det backward within its workspace budget, and the reviewed limits on a GPU (batch-local and packed indices outside the keys are empty slots). The two-GPU tests are marked multi_gpu. - test_prebuild_gpu.py: ROUTES over every call-time path (each head count from 4 to 128, 24 included; RoPE and NoPE; with and without a sink; thd, bshd and sbhd; no query rows; the indexer and its det top-k), each group checked after one prebuild. Co-Authored-By: Claude <noreply@anthropic.com>
Vendor the five packages of the gdn op into miles_ext_kernels/cake/_vendor/ with tools/vendor_cake.py, one level deep, recorded under ops.gdn of _vendor/VENDORED.json: - cake_gdn_det_train, cake_gdn_det_train_kcp, cake_gdn_train and cake_gdn_train_kcp: the FlashInfer Cake packages of flashinfer-ai/flashinfer eb3eb68d7deebbe619e05d93ecfd9796e4eaf2e3, byte for byte (85, 59, 80 and 67 files). This commit carries the fixes of the context-parallel packages, flashinfer-ai/flashinfer#8585 and #8623. - cake_gdn_fla_cp_det: the det kernels with FLA's context-parallel state passing generated for radixark/miles#3597, at 9d7543812b676a5b50dde107518782933b804b48: registry.json and the 56 kernel and launcher sources of miles/kernels/attention/gdn/cake/csrc. The tool adds an empty __init__.py, FlashInfer's loader cake_build.py, the TVM-FFI host surface csrc/cake_host_shim.h in place of the torch-extension one (excluded, its blob recorded) and a build.json generated from the registry, so these sources build with nvcc and load with tvm-ffi like the FlashInfer packages; sm_107a compiles the sm_103a records. The vendored files have the blobs the gdn packages, kernel sources and registry had in the previous layout. Commands, from the repository root, with a FlashInfer checkout at ../flashinfer and a miles checkout at ../miles: python tools/vendor_cake.py fi --checkout ../flashinfer \ --commit eb3eb68d7deebbe619e05d93ecfd9796e4eaf2e3 --op gdn \ --package flashinfer/experimental/cake_gdn_det_train \ --package flashinfer/experimental/cake_gdn_det_train_kcp \ --package flashinfer/experimental/cake/cake_gdn_train \ --package flashinfer/experimental/cake_gdn_train_kcp python tools/vendor_cake.py miles-pr --checkout ../miles --pr 3597 \ --commit 9d7543812b676a5b50dde107518782933b804b48 --op gdn \ --name cake_gdn_fla_cp_det --path miles/kernels/attention/gdn/cake \ --include csrc --include registry.json \ --exclude csrc/cake_host_shim.h --arch-from sm_107a=sm_103a Co-Authored-By: Claude <noreply@anthropic.com>
Lay out one namespace per kernel source, miles_ext_kernels.cake, with one shared runtime instead of a copy per package: - _runtime/spec.py: OpSpec and Flag, the whole declaration of an op (variants, the variants miles may choose, the miles selector, vendored packages, lazy entry points, options, its own argument checks). - _runtime/registry.py: OPS in layer order (empty here), op_spec, prebuild(op, **model_params) and release(op=None); prebuild remembers its paths per op, CUDA device (index and architecture), toolkit (nvcc and its CUDA version), cache directory, model parameters and op options until release(). - _runtime/config.py: the one argument surface (add_arguments, validate_args, configure) and the variant each op runs. An unconfigured process runs det; validate_args keeps miles' command-line default fast for a selected op and writes the resolved variant back, also through a hidden flag when miles may choose only one variant, so it reaches every process; deterministic mode resolves det, rejects an explicit fast and rejects an op without det, also for direct configure callers. configure replaces the whole configuration and the implementations read the variant when a call runs; that read looks the op's declaration up in sys.modules instead of importing it again, and variant and option, which run at every call, remember what they found per configuration and op registry. - _runtime/loader.py: the nvcc + tvm-ffi loader of the vendored packages, unchanged in behavior, plus install_op, load_package, prebuild_libraries, no_build (a build inside it raises, for the prebuild tests), device_arch, and a finder that resolves FlashInfer's ..cake.<name> imports to the one copy of _vendor.<name>, name by name. - _runtime/workspace.py: one workspace pool per process, device and stream, bounded by --cake-workspace-pool-bytes on every request, LRU eviction, release() and stats(). The default bound is 0: nothing is kept and a request is a plain torch.empty on the current stream (measured GDN training steps were no faster with a 1 GiB pool, which kept up to 0.75 GiB per device); a positive bound reuses buffers. A stream that captures a CUDA graph allocates outside the pool; a lock keeps the pool consistent across threads. take runs at every call of an op, so its host work is small: once a configuration's bound was read as 0 and no device holds a buffer, take counts the request and calls torch.empty on the resolved device; with a positive bound it reads the current device, its stream and their capture state with one CUDA runtime query each (no device switch for the current device, no stream object), each device keeps a running byte total, and a hit evicts nothing unless the bound was lowered. - miles_ext_kernels: SOURCES and the four hooks miles calls (add_arguments, validate_args, configure, release), torch-free and importing no op. Tests under tests/cake/_runtime/ run every check on a synthetic op, so they do not pass vacuously while no op is registered; tests/cake/helpers.py holds the shared accuracy and bitwise gates of the op tests, check_prebuild, which runs every route an op's test_prebuild_gpu.py declares (ROUTES) with every variant after one prebuild, with no build allowed, and POOLED, the 1 GiB bound a test of pooled buffers configures (the fixture pooled applies it). The op contract also forbids call-time switches as prebuild parameters and checks that importing an implementation rebinds every FlashInfer hook its vendored packages are allowed to keep. tests/cake/_runtime/fake_cuda.py is a fake CUDA device layer on which the pool's CUDA path (current device, streams, capture) runs on the CPU; bench_workspace.py next to it times take and the per-call lookups (a script, not a test). Co-Authored-By: Claude <noreply@anthropic.com>
tools/vendor_cake.py is the one way sources enter miles_ext_kernels/cake/_vendor/, one level deep, recorded in one _vendor/VENDORED.json grouped by op, with the source repository and commit of every package: - fi: FlashInfer Cake packages, byte for byte. - miles-pr: the generated kernel sources of a miles pull request (registry and csrc). The tool adds FlashInfer's loader, the TVM-FFI host surface tools/templates/cake_host_shim.h in place of the excluded torch-extension one, and a cake-build/1 build.json generated from the registry, so they build through the same loader as every other package instead of torch.utils.cpp_extension. - remove, and --check [OP ...]: the exact file set and blob of every package, each build.json, one FlashInfer commit per op, that a vendored module imports packages of its own op only, and no environment read or flashinfer import in vendored modules unless allowed (a flashinfer import only inside a listed function, a hook the op rebinds). --check --upstream compares the recorded files with the upstream commits (trees only) for both kinds of source. - merge: the git merge driver of VENDORED.json (.gitattributes), a three-way merge package by package, so layers that each add or change their own op merge and rebase without conflicts. Co-Authored-By: Claude <noreply@anthropic.com>
- pyproject.toml: depend on torch and apache-tvm-ffi>=0.1.11,<0.2; a test extra; the multi_gpu marker. - CI: the CPU tests with CPU torch, tools/vendor_cake.py --check, and a separate --check --upstream job. - pre-commit: leave the vendored tree and FlashInfer's loader template byte-identical. - README: install, the hooks miles calls, the Ops and Prebuild tables the op layers fill, variants and the scope of det, the prebuild contract, its routes and its limits, the workspace pool and the buffers packages keep outside it, the layout and the modules an op owns, vendoring, adding an op, the merge driver of the vendoring manifest, and how tests configure the pool. Co-Authored-By: Claude <noreply@anthropic.com>
Move the gdn op onto the cake layout. miles_ext_kernels/cake/gdn.py declares it (variants det and fast, --linear-attention-backend cake, --cake-gdn-recompute-state; importing it imports neither torch nor the kernels); miles_ext_kernels/cake/_impl/gdn/ holds the entry points; the five vendored packages load through the runtime loader. A call reads the variant and --cake-gdn-recompute-state when it runs: - no cp_context: det cake_gdn_det_train, fast cake_gdn_train; - a KcpContext: cake_gdn_det_train_kcp, cake_gdn_train_kcp; - an FLA cp_context, or a KcpContext call the KCP packages do not serve: cake_gdn_fla_cp_det, the det kernels of FLA's context parallelism, under either variant (fast logs it once). cake_gdn_fla_cp_det builds with nvcc through the runtime loader and loads with tvm-ffi (native.py); the torch-extension build is gone. prebuild() builds every library of the five packages for the device's architecture; a call no longer builds them on its first use. The third-party Triton kernels that a caller-supplied FLA cp_context runs (FLA's own context-parallel pre-processing) are not prewarmed: they compile in FLA on their first call. Behaviour changes: - --cake-gdn-variant fast under --deterministic-mode is now rejected, and so is configure(..., deterministic_mode=True) with fast: fast does not promise its bits, so deterministic mode runs det. Before, fast was accepted there. - Under any cp_context the contract is FLA's: initial_state and output_final_state raise ValueError when the call is made, and the context's cu_seqlens / cu_seqlens_cpu replace the arguments. Before, a KcpContext call with a final state ran the FLA-CP kernels, and a cu_seqlens_cpu other than the context's raised. - The scratch buffers of the det and fast packages and of cake_gdn_fla_cp_det go through the runtime's workspace pool instead of per-package pools kept until the process ends: at the default --cake-workspace-pool-bytes of 0 each is a fresh allocation per call; under a positive bound the pool keeps one buffer per role and stream, freed by release(). The 0-element state placeholder of a launch without an initial or final state holds no memory: it is one tensor per device, shape and dtype, not an allocation per call. Review fixes: - the FLA-CP state-scan planner never cuts a segment longer than the 128 chunks one pre-scan CTA covers: a long sequence that cannot be split within that keeps the serial scan instead of failing; - the FLA-CP path checks the extents of every input before a launch, as the det package does; - a call without a token returns an empty output and passes the state through, under either variant, also with an FLA context without a process group; - the autograd function of the FLA-CP kernels no longer materialises a missing incoming gradient: an unused final state adds no final-state gradient (the bits of the call without one), and a loss on the final state alone runs with a zero output gradient; - the scope of det's bits is documented (gdn.py, ops.py). The gdn extra installs flash-linear-attention>=0.5.2: the FLA-CP path calls the fla.ops.cp helpers of that release (with their state_v_first keyword) and is checked against it; 0.5.0's helpers lack the keyword, and the helpers' module of 0.5.1 differs from 0.5.2's. Co-Authored-By: Claude <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack 8/13 · base: main
Summary
This PR adds
miles_ext_kernels.cake.gqa, which implements packed (varlen) causal grouped-query attention with head dim 256 on the shared cake runtime. It covers the forward and the backward. It vendors one package from FlashInferc5db8d3,cake_gqa_train, and has one variant,fast. FlashInfer ships no det package for this kernel.In miles,
--full-attention-backend cakeselects the op (radixark/miles#4683). There is no variant flag. Passing--deterministic-modetogether with the cake backend is rejected at startup.Changes from the previous version of this layer:
max_seqlen_kis now a required host int. Its contract is documented, and an opt-in check (--cake-gqa-check-seqlens) enforces it.cake.prebuild("gqa")builds the four libraries the op launches.Stack
This is layer 8 of 13, on top of layer 7 (
mhc). The stack was rewritten as[Cake 0/13] layoutplus 13 op layers of three commits each (vendor / op / tests). Every PR targets main, so this diff contains layers 0 to 8. The previous version of this layer was reviewed in #8 ([Cake 8/11] gqa: packed causal GQA attention, head dim 256).This layer's commits:
06bcb00vendor: gqa packages438cc96gqa: packed causal GQA attention, head dim 256, on the cake runtime188f379tests: gqaChanges
From the layer 7 head to the layer 8 head: 34 files, +13,268 / −0. Per commit:
06bcb0022 files +11,863;438cc965 files +325;188f3797 files +1,080.cake/_vendor/cake_gqa_train/cake/_vendor/VENDORED.json(ops.gqablock)cake/gqa.py,cake/_impl/gqa/{__init__,ops}.pytests/cake/gqa/There are no
pyproject.tomlor tooling changes.Package
build.jsoncake_gqa_trainflashinfer/experimental/cake_gqa_train(21 files, Cake revision7ab5de0c)fwd,bwd_pre,bwd_fp32,bwd_rowmap,bwd_bf16,bwd_post) × {sm_100a, sm_103a, sm_107a}_vendor/. It needs no allow entry, because the vendored modules read no environment variable and do not import flashinfer. It has no local files.fwd,bwd_pre,bwd_bf16andbwd_post. These are the forward and the BF16 dK / dV backward.bwd_fp32andbwd_rowmapare not used.API
OPisname="gqa",variants=("fast",),selected_by=(("full_attention_backend", "cake"),),packages=("cake_gqa_train",). The module is torch-free and exportscheck_supported,is_supportedandvarlen_gqa_attention._impl.gqare-exports the same objects.varlen_gqa_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, *, max_seqlen_k, softmax_scale=None):q [T_q, H_q, 256]andk/v [T_k, H_kv, 256], withH_q / H_kveven.cu_seqlens_q/cu_seqlens_kare int32 device tensors withcu[0] = 0. They are non-decreasing and end atT_q/T_k.softmax_scaledefaults to256 ** -0.5.ValueErrorfor unsupported operands and in a process configured withdeterministic_mode. It raisesRuntimeErroron another GPU (vendored guard).max_seqlen_kis a hard contract. It is required, keyword-only and a hostint(notbool), and it must be at least every segment's key count.T_k, is correct but launches CTA pairs that only retire.TypeError. A negative value raisesValueError.packed_seq_params.max_seqlen_kv.--cake-gqa-check-seqlens(store_true, off by default;configure(cake_gqa_check_seqlens=True)) adds a check at every call. On the host and before any launch, it checks thatcu_seqlens_q/cu_seqlens_kstart at 0, never decrease and end at the token counts, and that no segment has more keys thanmax_seqlen_k. It costs one device synchronization per call and requires--full-attention-backend cake.check_supported(num_heads_q, num_heads_kv, head_dim=256)raisesValueErrornaming what the kernels do not run.is_supported(...)returns a bool. Both are torch-free.OP.validateruns invalidate_argswhen the op is selected and checks the heads of one tensor-parallel rank:num_attention_heads / TPquery heads overmax(num_query_groups, TP) / TPkey/value heads. The head dim comes from--kv-channelsorhidden_size / num_attention_heads.miles_ext_kernels.gqa:SUPPORTED_CAPABILITIES→cake.ARCH_BY_CAPABILITY,device_arch→cake.device_arch,prebuild→cake.prebuild("gqa").sm_100a,sm_103a,sm_107a; no prebuild parameters; tests intests/cake/gqa/.Variant routing
The op has one variant,
fast. It has no det variant and no KCP variant, so there is no variant flag to register.configure()without a gqa valuefast--full-attention-backend cake, not deterministicfast(no flag to write back)--deterministic-mode(flat or typed config) with--full-attention-backend cakevalidate_argsraisesValueErrorconfigure(..., deterministic_mode=True)ValueErrorwhen it runs; a laterconfigure()lets the next call throughOutput, LSE, dK and dV repeat bit for bit from run to run, for any valid
max_seqlen_k. dQ is summed with FP32 atomic adds in scheduling order, so it is not bitwise reproducible.Prebuild
cake.prebuild("gqa", *, device=None, max_workers=None)returns the paths of the 4 libraries (fwd,bwd_pre,bwd_bf16,bwd_post) for the device's architecture. It takes no model parameter. It raisesNotImplementedErroron a device other than sm_100a / sm_103a / sm_107a. It raisesRuntimeErrorwhen the toolkit is older than the architecture's minimum (CUDA 13.4 for sm_107a). The README Prebuild row reads:gqa | none: the four libraries of the forward and the BF16 dK / dV backward serve every head layout and length.Memory
("gqa", "bwd_scratch")on the current stream. The scratch is FP32[H_q/2, rows_pad, 256]plus two[H_q/2, rows_pad], withrows_pad = round_up(2*T_q + 128*B, 128), as views of one uint8 buffer.--cake-workspace-pool-bytesof 0 (layer 0), the pool keeps nothing. Every backward's scratch is then a plain allocation on the current stream, andrelease()frees 0 bytes.miles_ext_kernels.release()frees it.uint32tile counter (torch.zeros(1)). The counter is not pooled.(q, k, v, out, lse, cu_seqlens_q, cu_seqlens_k)and scalars. There is no device synchronization in steady state, and one per call when the check option is on._impl.gqa.release()has nothing of its own to drop.Review items
These items come from the review of the previous version: #8 (comment). The measurements are the reviewer's, taken on GB300.
max_seqlen_kwas accepted silently and gave wrong gradients. With 2048-key segments, bound 1024 gave dq rel-L2 0.20, dk 3.11 and dv 0.17, and bound 1920 gave dk 1.10. The output was unaffected. The reviewer asked for a hard contract and an opt-in debug check.max_seqlen_kis now a required, keyword-only host int. It is documented as a hard contract in the op module, the_impldocstring and the README Ops row. The opt-in--cake-gqa-check-seqlenschecks the boundaries and the bound on the host before any launch, at one sync per call. It is off by default.test_gqa_host.py::test_max_seqlen_k_is_a_required_host_int,::test_the_check_accepts_the_contract,::test_the_check_rejects_a_max_seqlen_k_below_the_longest_segment[0/1024/1920/LONGEST-1],::test_the_check_rejects_malformed_boundaries,::test_the_option_checks_before_any_launch_and_is_off_by_default;test_gqa_cpu_plan.py::test_the_backward_grid_covers_the_key_blocks_of_max_seqlen_k_only;test_arguments.py::test_the_check_option_is_off_by_default_and_needs_the_backend. GPU:test_gqa_gpu.py::test_the_check_rejects_a_max_seqlen_k_below_the_longest_segment[1024/1920/2047],::test_the_check_passes_a_valid_call_and_changes_no_bit,::test_steady_state_does_not_synchronize_and_the_check_doesmax_seqlen_k=Nonesized the backward grid by T. At 16 × 2048 packed, the backward took 2.20 vs 1.83 ms at H16/2 and 1.29 vs 1.09 ms at H8/1. out / dk / dv were bit-equal.Nonebound is removed. The caller passes the longest segment, and miles passespacked_seq_params.max_seqlen_kv.T_kis still allowed explicitly. Output, LSE, dK and dV do not depend on the bound given. The timings were not re-measured in this layer.test_gqa_host.py::test_max_seqlen_k_is_a_required_host_int. GPU:test_gqa_gpu.py::test_any_valid_max_seqlen_k_gives_the_same_out_lse_dk_dvvalidate_argsaccepted 24q / 4kv at TP 8, and the per-rank layouts 3/1 and 1/1 failed only at the first forward.check_supported/is_supportedare now exported, andOP.validatechecks the per-rank heads at startup.test_arguments.py::test_the_predicate_accepts_an_even_number_of_query_heads_per_key_head,::test_the_predicate_names_what_the_kernels_do_not_run,::test_validate_args_checks_the_heads_of_one_tensor_parallel_rank,::test_validate_args_leaves_heads_it_cannot_see_to_the_layer. GPU:test_gqa_gpu.py::test_rejects_unsupported_operandsvalidate_argsandconfigure).test_arguments.py::test_deterministic_mode_rejects_the_op,::test_the_typed_configuration_s_deterministic_mode_rejects_it_too,::test_configure_holds_the_deterministic_rules_and_the_option,::test_an_unconfigured_process_runs_fast;test_gqa_host.py::test_a_process_configured_for_deterministic_mode_rejects_the_call_when_it_runs. GPU:test_gqa_gpu.py::test_out_lse_dk_dv_repeat_bit_for_bit_in_one_process,::test_out_lse_dk_dv_are_the_bits_of_a_new_processCross-stack items from the review of #1 that touch this layer:
fast.cake.prebuild("gqa")builds the libraries, andtest_prebuild_gpu.pyrunsROUTESunderno_build().OP.flagsandOP.validate, and the runtime ownsadd_arguments,validate_argsandconfigure.test_gqa_gpu.py::test_default_and_side_stream_calls_agree_bit_for_bit,::test_the_backward_scratch_comes_from_the_pool_and_release_frees_it,::test_a_pool_bound_of_zero_keeps_nothing_and_changes_no_bit;test_gqa_host.py::test_at_the_default_bound_of_zero_the_backward_scratch_is_not_kept_and_release_frees_nothing.Tests
All tests are in
tests/cake/gqa/. Tests that need the pool configure the 1 GiBPOOLEDbound explicitly.CPU, torch-free:
test_arguments.py(13 tests) andtest_gqa_cpu_plan.py(5).CPU, with torch:
test_gqa_host.py(10).test_at_the_default_bound_of_zero_the_backward_scratch_is_not_kept_and_release_frees_nothing.test_the_backward_scratch_is_one_pooled_buffer_in_the_vendored_layoutnow uses thepooledfixture.GPU,
test_gqa_gpu.py(14):POOLED, then 0);test_two_gpus_give_the_same_bits(multi_gpu).GPU,
test_prebuild_gpu.py(3): one prebuild, then 15ROUTES, each run forward + backward underno_build():miles_exact_max_seqlen_k(segments 1000 / 517 / 2049 / 18, 8q / 2kv) andmax_seqlen_k_total_keys(bound =T_k);group_2,group_6andgroup_8;keys_fewer_than_queries,empty_segments,long_keysandsingle_segment;no_queriesandno_keys;softmax_scale(0.05),check_seqlens(option on) andside_stream;pooled(the workspace pool on with the 1 GiB test bound). It replacespool_off, since the default bound is now 0.The other two tests check that the prebuild returns the four libraries once built, and that a fresh process loads only the prebuilt libraries and never imports flashinfer.
The accuracy budget per tensor in
test_gqa_gpu.pyismin(1.10 × FlashInfer's FA-reference rel-L2, 2 × measured), checked byassert_rel_rmswithdetect_scale=1.02.{"out": 2.3e-3, "lse": 1.5e-7, "dq": 2.6e-3, "dk": 2.6e-3, "dv": 2.4e-3}.Validation
CPU, at each commit
Each commit was checked on a clean checkout. The image was
radixark/miles:dev-202611032236(Python 3.12.16, torch 2.14.1+cu130, apache-tvm-ffi 0.1.11), on a CPU-only x86_64 node with CUDA hidden (CUDA_VISIBLE_DEVICES=).pip install --no-deps -e . && pip install pytest, thenpython -m pytest -rfEs -p no:cacheprovider tests.PYTHONPATH=. python -m pytest -rfEs -p no:cacheprovider tests.python tools/vendor_cake.py --check,python -m tools.vendor_cake --checkandpre-commit run --all-files --show-diff-on-failure.python tools/vendor_cake.py --check gqa --upstreamandpytest -rs -m 'not multi_gpu' tests/cake/gqa.06bcb00(vendor)--checkrc 0; pre-commit rc 0, no difftest_vendored.py::test_every_vendored_op_is_registered; the op is registered in the next commit, as the per-commit rule allows438cc96(op)--checkrc 0; pre-commit rc 0, no difftest_op_contract.py::test_the_op_has_cpu_tests_and_a_prebuild_test[gqa]; the tests come in the next commit, as allowed188f379(tests)--checkrc 0; pre-commit rc 0, no diff;--check gqa --upstreamrc 0;tests/cake/gqa(image, CPU) 62 passed, 4 skippedThe skips are GPU tests on a node without CUDA, plus the torch-dependent tests in the torch-free venv.
vendor_cake.py --check --upstreamat the top of the stack (layer 13, with the sameops.gqablock and package tree) returned rc 0.GPU
The GPU tests ran on the top of the stack (layer 13). Its gqa files (
cake/gqa.py,_impl/gqa/,_vendor/cake_gqa_train/,tests/cake/gqa/) and itsops.gqablock are byte-identical to this layer's. Above this layer, only other ops' files change, plus the registry, README,VENDORED.jsonandpyproject.tomllines of later ops.The container was
radixark/miles:dev-202611032236(amd64 on B200, arm64 on GB300), with Python 3.12.16, torch 2.14.1+cu130, nvcc 13.0.88 (CUDA 13.0), apache-tvm-ffi 0.1.11 and pytest 9.2.0. Each op ran with its own JIT cache: "cold" means the cache started empty, and "warm" means the same cache was used again.Commands:
PYTHONPATH=. python -m pytest -rfEs -p no:cacheprovider tests/cake/gqa -m "not multi_gpu".-m multi_gpuon 2 GPUs.tests/cake/gqa/test_prebuild_gpu.pyalone on a new empty cache (B200), or as the first file of the cold run (GB300).test_two_gpus_give_the_same_bits)no_build()Gates on both B200 and GB300, with 0 failures:
fastguarantees the out / LSE / dK / dV bits.test_out_lse_dk_dv_repeat_bit_for_bit_in_one_process(4 miles head pairs) andtest_out_lse_dk_dv_are_the_bits_of_a_new_process(16/2, 8/1) passed.test_default_and_side_stream_calls_agree_bit_for_bitpassed, with the pool on and one scratch entry per stream.test_a_pool_bound_of_zero_keeps_nothing_and_changes_no_bitandtest_the_backward_scratch_comes_from_the_pool_and_release_frees_itpassed.release()returned 0.release()returned exactly the bytes held.prebuild("gqa")served every route underno_build(). A fresh process loaded only the 4 prebuilt libraries and never imported flashinfer.test_matches_the_fp64_referencepassed within the budgets above, on the FlashInfer suite layouts and the miles head pairs.Performance
In the copied FlashInfer text, "this PR" means the FlashInfer pull request that the text comes from, and "previous" means the Cake package on FlashInfer main before it, not this PR. A bare
#Ninside a quote or a table caption is flashinfer-ai/flashinfer#N.Package and variants
This layer vendors one package,
cake_gqa_train(21 files), from flashinfer-ai/flashinferc5db8d3ccd3fa3842812280f7f208718dc0b8610. Its package tree is02a458b65c52d3fbc3b65114ba79ce757456746c. The tree is the same in the earlier layer-08 revision of this stack, so the rows below are unchanged from that revision's performance section.c5db8d3or on FlashInfer main (367e4bc5c3679b5116db2e6bd914dd2e4e2d67ad). The package README states: "dQ and the row-map adds use atomic reductions, so they are not bitwise reproducible run to run."cake_gqa_trainfast). flashinfer-ai/flashinfer#7958: "The kernel has one data-type variant: BF16 inputs and outputs with FP32 accumulation. Its backward writes dK/dV in one of three forms: FP32 window gradients, an FP32 add into caller-owned destination rows through an int32 row map, or BF16." This op uses the BF16 dK/dV form.c5db8d3or on FlashInfer main (367e4bc5).FlashInfer pull requests of the vendored revision
c5db8d36bf67c4d7f1e016ce255051bb19d28b25002914802a458b65c52d3fbc3b65114ba79ce757456746c(= vendored)c5db8d3; no later commit on FlashInfer main changes it (checked at367e4bc5, 2026-11-05T07:50:28Z)509d751bcfb4a3bf0f6248e0856ace7c31d197dccdeca13795b41405d9d3c4ffa67bf80a3f1a5e4e(previous revision)flashinfer-ai/flashinfer#7958 delivered exactly the vendored revision. flashinfer-ai/flashinfer#7958: "In the backward main kernel of all three dK/dV forms, each thread that reads accumulators from tensor memory now releases them itself instead of handing them to a forwarder warp. The other kernels and the API are unchanged."
Baselines
index_add_; that table is copied below as context only. No body measures the vendored revision against a non-Cake kernel, and no chained ratio is computed here.Per-architecture coverage
build.jsonmin_cudainbuild.jsonThe package has no other targets. flashinfer-ai/flashinfer#7958: "Cake has no other data-type or layout variant of this kernel, and the kernel is prebuilt for
sm_100a,sm_103aandsm_107a." The package README states: "SM100 (B200 / GB200), SM103 (B300 / GB300) and SM107 (R200) are the supported architectures; another device raisesRuntimeError."sm_100a (GB200): the vendored revision against the previous Cake version
From flashinfer-ai/flashinfer#7958. On the same BF16 inputs, O, LSE, dK and dV are bitwise equal to the baseline on every shape, and dQ differs only by atomic-add order. On GB200, the backward geometric-mean speedup over the 9 shapes below is 1.023x (range 1.015x-1.036x), and the training step speeds up by 1.012x (range 1.006x-1.019x). Every shape is faster in both columns.
issue1-issue3are the three packed layouts of flashinfer-ai/flashinfer#7073 (context-parallel zigzag, CP 8 rank 0).zz_cp{N}_r{R}_{T}is rank R of an N-way context-parallel zigzag split of T tokens packed from several documents of irregular lengths (zz_cp1is a plain packed batch).The table times the FP32 window-gradient dK/dV form through the row map. This op uses the BF16 dK/dV form, which no FlashInfer body times. The change of flashinfer-ai/flashinfer#7958 is in the backward main kernel of all three dK/dV forms.
NVIDIA GB200 (sm_100a), CUDA 13.4 (prebuilt kernels): one training step of one attention layer under non-reentrant activation checkpointing (2 x forward + backward, including initialization of the FP32 dK/dV destinations), FP32 window gradients through the row map; "Backward" = preprocess + main + dQ postprocess; each value is the median over 24 rounds (4 runs, arm order alternated every round), each round the median of 10 back-to-back steps; this PR vs the previous Cake version on the same BF16 inputs; H_q = 24, H_kv = 4, head dim 256; speedup = previous ms / this PR ms.
sm_103a (B300) and sm_107a (Rubin)
From flashinfer-ai/flashinfer#7958: "With the same protocol at 10 rounds each, the backward geometric-mean speedup is 1.029x on NVIDIA B300 (sm_103a) and 1.055x on Rubin (sm_107a), and the training step speeds up by 1.018x and 1.033x."
No per-shape rows or absolute times exist for these targets.
Context only: the previous revision against FlashAttention-4
This table measures the first Cake version (flashinfer-ai/flashinfer#7774), not the vendored revision. flashinfer-ai/flashinfer#7958 states that the vendored revision gives O, LSE, dK and dV bitwise equal to that version. The two bodies use different timing protocols (a median of 10 steps against a median over 24 rounds of 10-step medians), so times are not compared across the two tables here. From flashinfer-ai/flashinfer#7774: "The baseline is FlashAttention-4 (CuTe DSL) HD256 varlen forward and backward (flash-attn-4 4.0.0b38.dev9 built from FlashAttention PR 2968, head f4015896), plus the FP32
index_add_workaround from flashinfer-ai/flashinfer#7073 that adds its BF16 window gradients into the destination rows. Against that baseline, the geometric-mean speedup over the 9 shapes below is 1.029x. Total time improves by 1.020x, and every shape is at least 1.003x faster."NVIDIA GB200 (sm_100a), CUDA 13.4: one training step of one attention layer under non-reentrant activation checkpointing (2 x forward + backward, including initialization), median of 10 steps; Cake vs FlashAttention-4 + FP32
index_add_on the same BF16 inputs; H_q = 24, H_kv = 4, head dim 256.Facts that FlashInfer does not report
Other facts in the vendored package
max_seqlen_k: the package README states thatmax_seqlen_ksizes the backward grid, and that "without it the bound isT_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."Notes
cake_gqa_trainships sm_107a libraries for all 6 stages and needs CUDA 13.4 to build them.prebuildrefuses an older toolkit withRuntimeError. No GPU run on sm_107a has been recorded for this tree yet.max_seqlen_kcontract is checked only with--cake-gqa-check-seqlens. A bound above the longest segment is correct but launches CTA pairs that only retire.🤖 Generated with Claude Code