Repository navigation
[Cake 13/13] mok: the Cake MoK expert-parallel MoE training layer - #14
Merged
Merged
Conversation
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>
The two Cake MoE grouped GEMM packages of FlashInfer
c5db8d3ccd3fa3842812280f7f208718dc0b8610, byte for byte, one level deep
under miles_ext_kernels/cake/_vendor/ and recorded under
ops.moe_grouped_gemm of _vendor/VENDORED.json:
cake_moe_grouped_det_gemm det flashinfer/experimental/cake_moe_grouped_det_gemm
cake_moe_grouped_gemm fast flashinfer/experimental/cake/cake_moe_grouped_gemm
Every blob and both Cake revisions are the ones the previous layout
vendored from the same commit. No vendored module reads the environment
or imports flashinfer, so neither package has an allow entry.
Vendored with:
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op moe_grouped_gemm \
--package flashinfer/experimental/cake_moe_grouped_det_gemm \
--package flashinfer/experimental/cake/cake_moe_grouped_gemm
Co-Authored-By: Claude <noreply@anthropic.com>
Vendor FlashInfer's chunked LM-head training packages at c5db8d3, one per
variant, into miles_ext_kernels/cake/_vendor/, with the lm_head block of
VENDORED.json:
- cake_lm_head_det_loss (flashinfer/experimental/cake_lm_head_det_loss,
102 files): det
- cake_lm_head_loss (flashinfer/experimental/cake/cake_lm_head_loss,
126 files): fast
Both package trees equal the upstream directories and the trees this op
shipped before (the same FlashInfer commit), blob for blob. Command, run
from the repository root with a FlashInfer checkout at ../flashinfer:
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op lm_head \
--package flashinfer/experimental/cake_lm_head_det_loss \
--package flashinfer/experimental/cake/cake_lm_head_loss
Co-Authored-By: Claude <noreply@anthropic.com>
Vendor the FlashInfer Cake package
flashinfer/experimental/cake_mok_bf16
of flashinfer-ai/flashinfer c5db8d3ccd3fa3842812280f7f208718dc0b8610
byte for byte into miles_ext_kernels/cake/_vendor/cake_mok_bf16/ (86
files, the package tree of that commit) and record it as ops.mok in
_vendor/VENDORED.json:
python tools/vendor_cake.py fi --checkout ../flashinfer \
--commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 --op mok \
--package flashinfer/experimental/cake_mok_bf16
where ../flashinfer is a FlashInfer checkout that holds the commit.
The op stays on c5db8d3. FlashInfer changes the package after it in two
performance PRs only, #8604 (vectorized MXFP8 dispatch quantize in the
sm_107a recompute) and #8717 (a Gluon MXFP8 weight quantizer); both
state outputs bitwise identical to before, and neither fixes
correctness.
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>
miles_ext_kernels.cake.moe_grouped_gemm declares the op and imports
nothing but the runtime's spec: det and fast, selected by
--moe-grouped-gemm-backend cake (which requires --train-backend
torchtitan), the two vendored packages, and the entry points
grouped_mm, check_grouped_dims and max_grouped_rows. The implementation
lives in cake/_impl/moe_grouped_gemm/ and builds its libraries through
the runtime's loader.
Changes from the previous layout:
- grouped_mm reads the configured variant at every forward, and its
backward runs the package of its forward. The variant= argument, the
module-level VARIANT and the op's own configure are gone: miles
configures the process with miles_ext_kernels.configure(...). A
process that never configured the op runs det; miles' command line
still defaults to fast. A fast call under
torch.use_deterministic_algorithms(True) runs det and logs it once,
as before: the fast weight gradient raises torch's alert there.
- The 32-bit limits are checked before anything runs. The kernels
address each output with 32-bit element offsets: sum_m * N,
sum_m * K and E * N * K must stay at most 2**31 - 1. grouped_mm checks
all three at the forward call, so a call whose input gradient would
overflow fails there instead of in the backward after the forward ran
(at N = 4096, K = 7168, 300,000 rows used to pass the forward).
check_grouped_dims takes num_groups and max_rows, so the model build
can refuse a configuration whose most rows per call exceed the bound,
and max_grouped_rows(N, K) names the bound (299,593 rows at N = 4096,
K = 7168).
- prebuild("moe_grouped_gemm") builds both variants' programs for the
device (no variant argument, no model parameter) and, on sm_100a and
sm_103a, the input gradient's probe of the memory domains, which the
first probed call used to compile: after it no call compiles.
- release() drops the workspaces both packages keep per device and
stream, with the queue splits probed on them.
- An unsupported device raises NotImplementedError (was RuntimeError).
README: the op's rows in the Ops and Prebuild tables.
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>
tests/cake/moe_grouped_gemm/: - test_arguments.py (CPU, torch-free): the variant flag through miles' hooks, deterministic mode, the torchtitan requirement, configure, and an import without torch. - test_limits.py (CPU, torch): the 32-bit bounds of check_grouped_dims and max_grouped_rows, and grouped_mm refusing at the forward, on meta tensors, a call whose input gradient or weight gradient would overflow. - test_moe_grouped_gemm_gpu.py (GPU): every output and gradient against a per-group FP32 reference with relative-RMS budgets that fail at a 2 % error (assert_close(atol=rtol=1e-2) passed y scaled by 1.05); det repeating bit for bit in one process and in a new one; fast giving det's forward and input-gradient bits and its unordered weight-gradient tail within the budget; the variant read at each call; det under torch's deterministic algorithms; calls alternating between the default and a side stream; release(); the 32-bit row bound at N = 4096, K = 7168 (300,000 rows refused at the forward; at 299,593 rows the last 2,048 rows of the output and of the input gradient and the whole weight gradient within the budget against closed-form FP64 references); empty input; padding rows past offs[-1]; refused calls; the toolkit check; a call for a device other than the current one (multi_gpu). - test_prebuild_gpu.py (GPU): ROUTES (every layer shape, empty experts, no rows, uneven routing over a sweep of row counts, padding rows, the 512-wide k tile, the stage twins at many rows per expert, a side stream) run with every variant after one prebuild, with no nvcc and no inline tvm-ffi build; a variant configured after a call runs without a build; every program the planners bind is prebuilt. 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>
miles_ext_kernels/cake/lm_head.py declares the op (torch-free): variants det and fast, selected by --log-probs-backend cake, the packages cake_lm_head_det_loss and cake_lm_head_loss, the entry point log_probs and the option --cake-lm-head-chunk-size (default 4096). Its validate hook keeps the startup checks of the previous arguments module: the chunk size range, Megatron, TP 1, the policy and SFT losses, no entropy, BF16, no LoRA on the output layer, no MuP, H a multiple of 512 and V of 256. miles_ext_kernels/cake/_impl/lm_head/ reads the variant and the chunk size when a call runs (the vendored autograd function keeps the forward's package for the backward). prebuild(*, device, max_workers) builds every stage log_probs can launch on the device's architecture (every stage of both packages except their candidate-support support_* stages, which serve other entries), with and without a temperature; release() clears the binding caches (and their workspaces) of both packages. registry.OPS and the README Ops and Prebuild tables get their lm_head lines. Review items: - prebuild(temperature=False) left the four temperature stages (gemm_logits_*_tau, row_finalize_tau) to the first temperature call: prebuild now always builds them; lm_head takes no model parameters. - A label >= V gave a finite, plausible logp: every call now queues the device-side assertion max(labels) < V ahead of the kernels (no host synchronization). Every negative label, not only -100, marks an ignored row; the docstrings say so. - On SM103a with a CUDA 13.0 toolkit every forward waits for the device, also with num_valid_rows: the packages keep those calls on their host-count path (their mitigation of a ptxas 13.0 fault). It is documented, and the first call or prebuild on such a device logs it once; any other toolkit (12.9, 13.1 or newer) avoids it. - dW is BF16 only: it stays so. The log-probability entries of both packages return a BF16 dW, and the autograd engine casts a gradient to its input's dtype, so an FP32 option for a BF16 W would change nothing. The docstring states the FP32 [V, H] accumulator and the memory the backward holds. - V must be a multiple of 256, the kernels' logits tile: a vocabulary such as 151,936 is rejected at startup with that reason (padding it would change the softmax), and documented. Behavior changes against the previous lm_head package: - A process that never configured the op runs det (was fast); miles' CLI still resolves an unset --cake-lm-head-variant to fast, and to det under --deterministic-mode. - fast no longer turns into det while torch.use_deterministic_algorithms is on: the variant comes from configure alone, where deterministic mode rejects fast. - log_probs takes no variant= or chunk_size= argument; device_arch, prebuild and the supported capabilities come from miles_ext_kernels.cake, and an unsupported GPU raises NotImplementedError (was RuntimeError). - An X off CUDA raises ValueError (was RuntimeError). 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>
tests/cake/lm_head/: - test_arguments.py (CPU): the flags through miles' hooks, variant resolution with and without deterministic mode, the startup rejections (151,936 with the logits-tile reason) and the process configuration. - test_package.py (CPU torch): prebuild asks for every stage of both packages with and without a temperature on each architecture; only SM103a programs built by CUDA 13.0 take the host-count path, logged once; the program registries read the toolkit through the loader; release clears the binding caches and imports nothing. - test_lm_head_gpu.py: the accuracy gate against an FP32 oracle (checked to fail a 2 % error), det bit for bit in one process and in a new one, fast's det bits, empty calls, frozen inputs, given counts, zero-gradient rows, the configured variant read at call time, a side stream, a second GPU (multi_gpu), a given count that does not wait for the device behind a queued sleep (and the SM103a CUDA 13.0 exception), negative labels, a label >= V ending the stream in a fresh process, dW as the BF16 cast of the FP32 accumulator, and the geometry and device errors. - test_prebuild_gpu.py: ROUTES over temperature off, scalar, per-row and 1.0, one row to several chunks, chunk sizes 1024 and 65535, given counts, ignored and zero-gradient rows, frozen inputs and no-grad forwards at two geometries; after one prebuild none compiles, for either variant. Co-Authored-By: Claude <noreply@anthropic.com>
Vendor FlashInfer's Cake package flashinfer/cake_gated_act_mxfp8_backend (fused SwiGLU + MXFP8 quantization, forward and backward) of flashinfer-ai/flashinfer c5db8d3ccd3fa3842812280f7f208718dc0b8610 into miles_ext_kernels/cake/_vendor/cake_gated_act_mxfp8_backend/, byte for byte (131 files, package tree ffec0b8; the tree is the same at FlashInfer 7a13d7ba9b75c9bb9c02452df36e3f34417b5d57), and record it as ops.gated_act_mxfp8 in _vendor/VENDORED.json, from the repository root with a FlashInfer checkout at ../flashinfer: python tools/vendor_cake.py fi --checkout ../flashinfer \ --commit c5db8d3ccd3fa3842812280f7f208718dc0b8610 \ --op gated_act_mxfp8 \ --package flashinfer/cake_gated_act_mxfp8_backend Its Python modules read no environment variable and import flashinfer only in the pinned cake_build.py loader, so the entry has no allow list. build.json (Cake revision ce6037e6) lists 24 libraries for sm_100a, 28 for sm_103a and 23 for sm_107a (min_cuda 12.8, 12.9 and 13.4). The op module registers the package in the next commit. Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.cake.gated_act_mxfp8 runs the vendored Cake fused SwiGLU + MXFP8 quantization (cake_gated_act_mxfp8_backend), forward and backward, on SM100a, SM103a and SM107a. The forward quantizes silu(gate) * up of a BF16 [M, 2K] gate / up input (M and K positive multiples of 128), the backward the gate and up gradients [M, 2K] of a BF16 grad_output [M, K]: E4M3 values with RCEIL E8M0 scales per 32 elements of a row and / or a column, the scales in the swizzled 128 x 4 layout of the MXFP8 GEMM operands. OP: one variant, det. A call runs one program of its route chosen by M * K alone, and every program of a route writes the same bytes, so the outputs are a function of the inputs and repeat bit for bit. No miles argument selects the op, so it registers no flag; it has no option. The exports are the package's own entry points, silu_and_mul_mxfp8_quantize and silu_and_mul_mxfp8_quantize_backward. _impl.gated_act_mxfp8 installs the runtime loader on the package. prebuild(device=None, max_workers=None) takes no model parameter (M and K are runtime arguments) and builds every library of the device's architecture: the six routes (forward and backward; rowwise, colwise or both), each with every program its M * K rule picks. An unsupported device raises NotImplementedError, a toolkit that cannot build the architecture RuntimeError (SM107a needs CUDA 13.4). release() has nothing to free: no workspace, plan or graph, and the loaded libraries stay loaded, since a captured CUDA graph may still launch them. registry.OPS gains the op; the README lists it in the Ops and Prebuild tables. Co-Authored-By: Claude <noreply@anthropic.com>
… not the README The README no longer lists ops in a table; the op module and implementation docstrings already state that neither variant is batch-invariant.
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 13/13 · base: main
Summary
Adds
miles_ext_kernels.cake.mok, the Cake MoK expert-parallel MoE training layer. It comes from the FlashInfer packagecake_mok_bf16, vendored byte for byte from flashinfer-ai/flashinfer@c5db8d3.The layer covers
E = ep_size × local_expertsrouted experts in BF16 or MXFP8, plus one BF16 shared expert, with SwiGLU or clamped SwiGLU. It provides forward, context recompute and backward.The module exports FlashInfer's three entry points, with FlashInfer's signatures and defaults:
prepare_mok_bf16(...)returns an adapter withbuild_schedule,forward,recompute_forward_contextandbackward.create_mok_bf16_workspace(...)returns a caller-owned(config, workspace). Every rank creates it collectively.quantize_mok_mxfp8_weights(gate, up, down)returns the MXFP8 forward and backward weights for the routed experts.Two behaviors are new in ext:
NoMulticastErrorat creation. Before, it failed in the first dispatch.prebuildcompiled then always runs.There is one package and one variant, det. The bits do not depend on the communication-SM split. FlashInfer has no fast (non-deterministic) and no KCP
cake_mok*package (checked atc5db8d3, at7a13d7baand at main79cf05572cab).No miles argument selects the op (
OP.selected_by = ()). It registers no flag and has no miles call site, so there is no paired miles PR.Stack
Layer 13 of 13, the top of the stack; this is a new PR. It builds on layer 12 (
gated_act_mxfp8, new PR).Every PR in the stack targets main, so this diff also contains layers 0-12. This layer is these three commits:
66ff458725aadf5cbaba2c53Changes
9f72261d..baba2c53: 103 files, +274,734 / -0.cake/_vendor/cake_mok_bf16/66ff4587cake/_vendor/VENDORED.json(ops.mok)66ff4587cake/mok.py(+31),cake/_impl/mok/__init__.py(+28),cake/_impl/mok/ops.py(+316)25aadf5c_runtime/registry.py(OPS),README.md(extras, Ops and Prebuild rows),pyproject.toml(mok = ["triton>=3.8"])25aadf5ctests/cake/mok/baba2c53This layer changes no tooling or CI files. Its only layer-0 edits are the
"mok",entry inOPS, the three README rows and themokextra inpyproject.toml.Package (
66ff4587)min_cudacake_mok_bf16flashinfer/experimental/cake_mok_bf16d2a4ce7aVendored with this command:
Cake revision:
433416df9e68bc09d99f392f96097445bc71a54f. There are noallowentries.sources.jsondefines 48 kernel roles per architecture:zero,expert_layout,gather,barrier;count_*/rows_*for the 9 scheduler layouts;pad_{1,4,8,16,32,64}andepilogue_{forward,backward}_{2,8};The package stays on
c5db8d3. The later upstream changes are perf-only, state bitwise-identical outputs, and fix no correctness issue:The last upstream change in the vendored tree was perf(cake_mok_bf16): vectorized shared-memory MXFP8 dispatch quantize in the sm_107a backward (#7821) flashinfer-ai/flashinfer#8546.
nvcc needed per architecture: sm_100a 12.8, sm_103a 12.9, sm_107a 13.4.
The vendored README says the MXFP8 forward and recompute kernels need CUDA 13.1. At this commit, however, every
cvt.rn.satfinite.e4m3x2.bf16x2in the MXFP8 sources sits under a CUDA 13.1 guard, with the FP32 form of the conversion for older nvcc. CUDA 13.0 therefore builds every SM100a and SM103a kernel; a CPU test checks the guard.Op (
25aadf5c)Module layout:
cake/mok.pyis torch-free; its entry points load_impl.mokon first access._impl/mokholds the model-keyedroles,librariesandprebuild, the multicast check (NoMulticastError), a weight quantizer that does not depend on the weight layout, andrelease.prepare_mok_bf16(*, ep_size=16, local_experts=16, topk=8, clamped_swiglu=False, fp32_wgrad=False):prepare(swiglu_limit, mxfp8)/prepare_recompute(...);ValueErrorfor a layout or top-k that is not served.create_mok_bf16_workspace(*, group, device, num_local_tokens, hidden_size, topk, source_capacity=None, fwd_num_comm_sms=40, bwd_num_comm_sms=40, minibatch_size=4096, macrobatch_size=32768, schedule_capacity_multiplier=3/16, precision="bf16", imbalance_comm_sms="auto", imbalance_threshold="auto"):NoMulticastError, naming the NVLink-domain requirement.NoMulticastErroris aRuntimeErrorsubclass defined inmiles_ext_kernels.cake._impl.mok.ops. Before, such a group failed in the first dispatch.quantize_mok_mxfp8_weights(gate, up, down):gate/upas[E_local, I, H]anddownas[E_local, H, I];prebuildcompiled always runs. The bytes do not depend on the layout.What the layer serves:
(ep_size, local_experts):The backward returns
dx, the score gradients and the six routed + shared weight gradients. They come back in BF16, or withfp32_wgradthey are added into six caller-owned FP32 accumulators.Runtime requirements:
torch.distributedNCCL group over the expert-parallel ranks, with its CUDA device current;mok,triton>=3.8) for the weight quantizer. Without it, PyTorch operations write the same bytes.Variant
configureaccepts det only. There is no no-op fast variant.Prebuild
cake.prebuild("mok", *, device=None, max_workers=None, ep_size=None, local_experts=None, topk=None, hidden_size=None, intermediate_size=None, mxfp8=True, clamped_swiglu=None, fp32_wgrad=None).ep_sizeandlocal_expertsare given together (one scheduler layout) or not at all.topk,clamped_swigluorfp32_wgradis unset, every value the kernels serve is built.mxfp8(default true) includes the MXFP8 routed kernels. It then requireshidden_sizeandintermediate_size, the routed weight shape the Triton quantizer is compiled for. Triton specializes on the widths, not on the expert count.topk,clamped_swigluandfp32_wgradall unset: all 48. Under the defaultmxfp8=Truethis needshidden_sizeandintermediate_size; leaving them out raisesValueError.topk+clamped_swiglu+fp32_wgradset: 15, or 12 withmxfp8=False.clamped_swiglu/fp32_wgradunset): 25.mxfp8, it also compiles the weight quantizer's Triton kernel on the device.prepare/prepare_recompute.ValueErrorfor parameters that are not served, checked before touching a device;NotImplementedErrorfor another GPU;RuntimeErrorfor a toolkit that is too old (SM107a needs CUDA 13.4).Memory
--cake-workspace-pool-bytes, default0) does not apply.create_mok_bf16_workspaceallocates symmetric-memory buffers sized for every rank'snum_local_tokens. Adapters, workspaces and captured graphs belong to the caller.release()drops what the vendored package caches for the process: its loaded kernel modules (jit._load) and its per-device constants (mxfp8._scale_factor). A new adapter then reloads the libraries from the cache without a build.Tests
All tests are in
tests/cake/mok/. The helpers are_cases.py,_reference.pyand_worker.py.CPU
test_arguments.py: the declaration and the variant rules (table below), plustest_the_entry_points_are_lazy.test_ops_cpu.py: the prebuild selection and the multicast check (table below), plus:test_the_runtime_has_what_the_package_uses,test_the_package_builds_through_the_runtime_loader;test_the_weight_quantizer_writes_the_recipe_bytes_on_cpu,test_weights_of_another_layout_are_copied_and_quantized_alike;test_release_drops_the_loaded_modules_and_the_device_constants.test_package.py:test_every_role_is_built_for_every_architecture,test_the_roles_are_exactly_the_manifest_libraries(3);test_no_library_needs_a_newer_toolkit_than_the_loader_names(3);One GPU
Setup: EP 1, layout (1, 4), top-2, hidden 512, intermediate 256, routing
randomandempty.test_mok_gpu.py::test_bf16_cases(10 = 5 cases × 2 routings) and::test_mxfp8_cases(8 = 4 × 2) run the 9 cases. Each case also checks:test_mok_gpu.py:test_outputs_repeat_bit_for_bit_in_a_fresh_process(2);test_the_weight_quantizer_writes_the_recipe_bytes_on_the_device;test_unsupported_layouts_raise;test_the_libraries_build_through_the_loader_without_flashinfer(2).Prebuild (
test_prebuild_gpu.py)ROUTESholds:bf16,bf16-clamped,bf16-recompute,bf16-fp32-wgrad,bf16-clamped-recompute-fp32-wgrad,mxfp8,mxfp8-fp32-wgrad,mxfp8-clamped-recompute,mxfp8-clamped-fp32-wgrad;tokens-0,tokens-1,tokens-777;new-workspace,other-comm-split,weight-quantizer-layouts.The check does not depend on what the process ran before:
release("mok")drops the loaded kernel modules, and the vendored Triton kernels' in-process caches are cleared, soprebuildmust compile the quantizer into a private, emptyTRITON_CACHE_DIR.no_build(), and the private Triton cache gains no entry.load_library, must equal exactly the onesprebuildreturned.The narrow prebuild (
clamped_swiglu=False,fp32_wgrad=False: 15 libraries, or 12 withmxfp8=False) gets the same check on that form's routes.Multi-GPU (
multi_gpu)test_mok_distributed_gpu.py::test_expert_parallel_training_on_one_noderuns 6 launches:ep4-random,ep4-empty,ep4-empty-mxfp8-split;ep8-random,ep8-empty,ep8-empty-mxfp8-split.All launches use top-8 and unequal source counts, with empty ranks. After a prebuild of the layout, every case runs on every rank. A second launch in fresh processes then repeats every output bit for bit.
emptyrouting sends every row to rank 0. On GPUs that have an imbalance split row (B200 for BF16 and MXFP8; SM107a with 212 SMs for MXFP8), each rank must log the switch.*-mxfp8-splitis skipped where the GPU has no MXFP8 row (GB300).NoMulticastError(no NVLink multicast), the test skips with that error. If only some ranks report it, the test fails.Cross-stack review items
There is no earlier PR for this op. The cross-stack items from the review of the earlier stack apply as follows. All of these tests passed on B200 and GB300.
test_arguments.py::test_every_mode_runs_det(2),::test_configure_accepts_det_only,::test_the_op_is_registered_with_one_det_variant_and_no_selectorprebuildcovers every call-time switch and warms the Triton quantizer; the GPU check does not depend on test ordertest_ops_cpu.py::test_with_nothing_set_the_libraries_are_the_build_manifest(3),::test_one_model_builds_its_layout_topk_form_and_weight_gradients(3),::test_an_unset_parameter_builds_every_value,::test_prebuild_rejects_parameters_the_kernels_do_not_serve(13); GPUtest_prebuild_gpu.py::test_after_prebuild_no_route_compiles,::test_after_a_prebuild_of_one_form_its_routes_compile_nothing[mxfp8/bf16],test_mok_gpu.py::test_prebuild_builds_the_model_libraries_and_new_token_counts_load_none[random/empty]test_arguments.py::test_no_argument_is_registeredassert_rel_rmsagainst an FP32 reference, at twice the error of an independent BF16 reference of the same layer measured in the same run. The gate must catch a 2 % scale. MXFP8 keeps FlashInfer's relative L2 bounds on the routed path, and FlashInfer's own gates stay. Clamped SwiGLU needs more care. BF16 rounding of the gate and up GEMM outputs moves elements across the limit, where the clamp's derivative jumps. As a result,dxand the routed and shared gate / up weight gradients are several % off FP32 in any BF16 implementation, including the BF16 reference, and a 2 % scale stays inside twice that error. These tensors keep that budget without the 2 % check. They are also gated, with the 2 % check, at twice the BF16 reference's error against an FP32 reference whose clamp decisions are made on its GEMM outputs rounded to BF16.test_mok_gpu.py::test_bf16_cases(10),::test_mxfp8_cases(8)test_a_call_with_no_token(2); ROUTEStokens-0test_calls_alternating_with_a_side_stream_repeat_the_bits[random/empty]NoMulticastError(RuntimeError) at creation. The multi-GPU test skips with that error where every rank reports it, and fails if only some ranks do.test_ops_cpu.py::test_a_multi_rank_workspace_needs_a_multicast_address; GPUtest_mok_distributed_gpu.py::test_expert_parallel_training_on_one_node[...]Validation
CPU, at this layer head and at each of its commits. Container
radixark/miles:dev-202611032236(Python 3.12.16, torch 2.14.1+cu130, tvm-ffi 0.1.11), with CUDA hidden.66ff458725aadf5cbaba2c53(head)pip install --no-deps -e . && pip install pytest;python -m pytest -rfEs -p no:cacheprovider testsPYTHONPATH=. python -m pytest -rfEs -p no:cacheprovider testspython tools/vendor_cake.py --checkandpython -m tools.vendor_cake --checkpre-commit run --all-filespython tools/vendor_cake.py --check mok --upstreampython tools/vendor_cake.py --check --upstreampython -m pytest -rs -m 'not multi_gpu' tests/cake/moktests/cake/_runtime/test_vendored.py::test_every_vendored_op_is_registeredafter the vendor commit, andtests/cake/_runtime/test_op_contract.py::test_the_op_has_cpu_tests_and_a_prebuild_test[mok]after the op commit. Any other failure fails the gate.9f72261d): +20 passed / +6 skipped in the venv, +48 / +3 in the container.GPU, at this head (
baba2c53). Both sites used theradixark/miles:dev-202611032236container (amd64 on B200, arm64 on GB300) with Python 3.12.16, torch 2.14.1+cu130, nvcc 13.0.88, apache-tvm-ffi 0.1.11, triton 3.8.0 and pytest 9.2.0.B200: 8× B200 in one node, sm_100a, 148 SMs, driver 580.82.07.
-m 'not multi_gpu', empty cache)test_prebuild_gpu.pyalone, new empty cache)multi_gpu(8 GPUs)ep4-random,ep4-empty,ep4-empty-mxfp8-split,ep8-random,ep8-empty,ep8-empty-mxfp8-split(every rank 9 cases, 0 failures, checked and fresh launches)GB300: 4× GB300 in one node, sm_103a, 152 SMs, driver 580.159.03.
test_prebuild_gpu.pyfirst)multi_gpu(4 GPUs)ep4-random,ep4-empty), 4 skippedThe GB300 skips come from hardware and feature limits:
ep4-empty-mxfp8-split: "no mxfp8 imbalance split on this GPU: the run with the bf16 workspace covers this routing". FlashInfer's imbalance split table has no MXFP8 row for 152-SM sm_103a.ep8-random,ep8-empty,ep8-empty-mxfp8-split: "needs 8 GPUs". The GB300 node has 4.Gates inside those runs:
test_outputs_repeat_bit_for_bit_in_a_fresh_process[random, empty], the per-case eager, graph, split and fresh-process repeats, and the multi-GPU second launch.test_calls_alternating_with_a_side_stream_repeat_the_bits[random, empty]) passed.Cold times: on B200, eight test runs shared the node's CPUs while compiling. On GB300, four ran concurrently (34 CPUs each). Cold times include pytest start-up and every test, not just a bare
prebuild.sm_107a: validation of this tree pending.
Performance
In the copied FlashInfer text, "this PR", "before", "previous" and "main" refer to that FlashInfer pull request and the package before it, not to this PR. A bare
#Ninside copied text is flashinfer-ai/flashinfer pull request or issue N.Packages and variants
This layer vendors one package from flashinfer-ai/flashinfer
c5db8d3ccd3fa3842812280f7f208718dc0b8610. The vendored package tree is identical to the FlashInfer tree at that commit.cake_mok_bf16(FlashInferflashinfer/experimental/cake_mok_bf16)d2a4ce7af694c00b0bad1514c01e8785cbef47dacake_mok*package exists underflashinfer/experimentalatc5db8d3.c5db8d3.This op is new in this stack, so no earlier facts are reused. Every copied table row was checked verbatim against the current FlashInfer PR bodies on 2026-11-05.
FlashInfer pull requests of the vendored revision
c5db8d3e32cf76ebb41636f4556d9675949e8a107891da5d2a4ce7a(= vendored)d2a4ce7a(= vendored)c5db8d35b2180ccc4534a97151f3b3149eb6d54baf5f9a1fc87da0e(previous revision)fc87da0e23d065c54c7d555a93d4a9c17627b439b4249bbf6e91f4bac0b7622c40f92526b5c6a513720ddb97f04a491462561d7324a4513624a451368d0fa1cea05117c1daefafb4f8529a374919751502498a7502498a751e47a692ef2a37e5c1decd7d43944d6e73623b345adad13c5adad13cbuild.jsonand the four sm_107a MXFP8 backward kernels (csrc/sm_107a/cake_mok_backward_mxfp8{,_clamped,_fp32,_clamped_fp32}.cu). Its body measures sm_107a only: "Before is FlashInfer main before this PR, which already includes perf(cake_mok_bf16): vectorized shared-memory MXFP8 dispatch quantize in the sm_107a forward (#7821) flashinfer-ai/flashinfer#8506."build.jsonlibrary entries for these targets atc5db8d3are identical to those after flashinfer-ai/flashinfer#7326. Every kernel source change after flashinfer-ai/flashinfer#7326 is undercsrc/sm_107a/. After flashinfer-ai/flashinfer#7326, host Python changed in flashinfer-ai/flashinfer#7870 (api.py,backend.py,mxfp8.py) and flashinfer-ai/flashinfer#8346 (mxfp8.py, the MXFP8 weight quantizer). These checks compare the package files at each merge commit through the GitHub API. flashinfer-ai/flashinfer#7326 is the latest FlashInfer body with B200 or GB300 numbers.Baselines
Per-architecture coverage
Library counts and
min_cudacome frombuild.jsonat the vendored revision. Itsarch_flagslist exactly sm_100a, sm_103a and sm_107a.min_cudaNo target is compile-only or unsupported. flashinfer-ai/flashinfer#8546: "every variant runs on sm_100a, sm_103a and sm_107a." No FlashInfer number exists for a 216-SM R200.
sm_107a: the vendored revision against the previous Cake version
From flashinfer-ai/flashinfer#8546, verbatim:
sm_107a: earlier steps, each against the package before it (context only)
The earlier sm_107a steps, flashinfer-ai/flashinfer#8506, flashinfer-ai/flashinfer#8346, flashinfer-ai/flashinfer#8302 and flashinfer-ai/flashinfer#7870, each measured their own head against FlashInfer main before it. None measured the vendored revision. Their tables are copied verbatim in the follow-up comment "Performance (continued)".
sm_100a (B200) and sm_103a (GB300)
No FlashInfer body measures the vendored revision on B200 or GB300. The latest measurement is flashinfer-ai/flashinfer#7326. Its kernel sources and build entries for these targets are the same as the vendored ones; host Python and the MXFP8 weight quantizer changed later (flashinfer-ai/flashinfer#7870, flashinfer-ai/flashinfer#8346). From flashinfer-ai/flashinfer#7326, verbatim:
Its four per-shape tables (B200 EP8, B200 EP4, GB300 EP4, GB300 EP8; 28, 28, 28 and 10 rows, BF16 and MXFP8) are not copied here, because their row labels include a model name. See flashinfer-ai/flashinfer#7326.
Earlier revisions (context only)
Each of these FlashInfer pull requests is an ancestor of
c5db8d3and reports performance against the package before it. None of them measured the vendored revision.5b2180ccc4534a97151f3b3149eb6d54baf5f9a123d065c54c7d555a93d4a9c17627b439b4249bbf40f92526b5c6a513720ddb97f04a491462561d738d0fa1cea05117c1daefafb4f8529a37491975151e47a692ef2a37e5c1decd7d43944d6e73623b344a35122d9e80d4aa1ef2fe3de2327d468510134009d482e1abbd833094c597c083fa8553040c09baaf8c006ba389efb9202230f9f3306af6424568bee44067b1fbfce66f547d48c41976b1a57366b8f8Later revisions on FlashInfer main (not vendored)
FlashInfer main has later revisions of this package that are not in
c5db8d3: flashinfer-ai/flashinfer#8604 (3ded62b988e2557f9313b4acd8bc22920f39a840) and flashinfer-ai/flashinfer#8717 (bcb0416a7c8ea610685c6c74bf3c3410f8369781). Their numbers do not describe the vendored package and are not used here.Missing facts
Notes
test_package.py::test_every_role_is_built_for_every_architecturechecks on CPU.prebuildraisesRuntimeError.NoMulticastError, and the multi-GPU test skips with it.ep4-empty-mxfp8-splitskips because there is no MXFP8 imbalance split row for 152-SM sm_103a. The EP 8 launches need 8 GPUs in one node and ran only on B200.prebuildrun, outside or before CUDA graph capture.hidden_sizeandintermediate_size.🤖 Generated with Claude Code