Repository navigation
[Cake 1/13] gdn: chunked gated delta rule, det and fast, with KCP and FLA-CP - #1
Conversation
Review of layer 1 (gdn + shared infra) — suggest request changesConfirmed on GPU
Cleared on GPU
Code-level
Across the stack
Verified on two GB300 nodes (sm_103a), CUDA 13.0, torch 2.14.1+cu130, at the stack head |
Repo layout — please standardize it before splitting into stacked PRsMore kernel sources will come later, so group everything by source under one namespace, with a single layout for every op: What this fixes at the stack head (
Suggested order: land a "[Cake 0/11] layout" PR first, then put the 11 layers on top of it. |
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>
tests/cake/gdn/, on the CPU: the arguments (fast rejected under deterministic mode), the five packages and their build manifests, the planner's 128-chunk segment cap, calls without a token, host boundaries that end before the call's tokens (rejected before any CUDA call), the state placeholder shared and the packages' other scratch taken from the workspace pool, and the KCP and FLA context routing on the packages' torch emulation (FLA's context-parallel contract, a call the KCP package does not serve logged once). On a GPU: det and fast against the FP64 recurrence, FLA and each other; det bit for bit run to run, in a new process, under PyTorch's deterministic algorithms and on a side stream alternating with the default one; the extents checked before a launch; one sequence of up to 128k tokens on the det package and on the det kernels with FLA's context parallelism; the scratch in the workspace pool under a positive bound (helpers.POOLED, which every pool test configures) at the bits of the default bound 0, which keeps no pooled byte after a training step; the det kernels with FLA's context parallelism from the loader's libraries, where an unused final state passes no gradient (the bits of the call without one) and a loss on the final state alone backpropagates; the KCP packages on one GPU and, marked multi_gpu, with context parallelism 2 and 4, where a final state is rejected on every rank when the call is made. A case that raises on every rank before any communication fails its own tests with the ranks' tracebacks, and the other cases of the shared multi-GPU run still report; a case that raises on some ranks only is not isolated, and the process group's 5-minute timeout bounds the resulting hang. test_prebuild_gpu.py declares ROUTES, one call per call-time path, and checks that nothing compiles after one prebuild, and that a call runs the variant configured when it is made. fast makes no promise about its bits: its tests compare within the accuracy budgets only. The recomputed det backward is compared within the budgets of the saved one (RECOMPUTE_SAME_BITS = False). The expected scan plans are the cost model's optima at the 148 and 152 SMs of the B200 and GB300 parts: the tests that check them on a device plan at those SM counts, since a part with another SM count plans other segment lengths. Where torch.profiler records no CUDA API call (one known kernel launch shows none either), the launch-count test skips and names the device. Co-Authored-By: Claude <noreply@anthropic.com>
c62be50 to
ad8eeef
Compare
yyihuang
changed the title
Performance (1/5) for [Cake 1/13] gdn, continued from the PR bodyPerformanceIn the copied FlashInfer text, "this PR", "new", "previous" and "main" refer to that FlashInfer pull request and to the FlashInfer package before it, not to this PR. Tables, table captions and quoted paragraphs are copied from the pull request bodies without change; a copied heading is shown in bold. Packages and variantsThis layer vendors four packages from flashinfer-ai/flashinfer
The earlier version of this PR vendored only Libraries in FlashInfer pull requests of the vendored revisionEvery merge commit below is an ancestor of
Baseline identities: the Not vendored: these FlashInfer main changes come after Baselines
Per-architecture coverage
None of the pull requests above reports B300 numbers. Deterministic, single-rank, sm_107a (exact revision)Source: flashinfer-ai/flashinfer#8695. Baseline: the det package on FlashInfer main before it. No non-Cake baseline.
NVIDIA R200 (sm_107a, 216 SMs), CUDA 13.4 (torch 2.14.1+cu130): BF16 inputs, head dim 128, chunk 64; forward + backward (f+b) and forward (fwd) per call, ms (median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR.
NVIDIA R200 (sm_107a, 212 SMs), CUDA 13.4 (torch 2.14.1+cu130): BF16 inputs, head dim 128, chunk 64; forward + backward (f+b) and forward (fwd) per call, ms (median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR.
NVIDIA R200 (sm_107a, 208 SMs), CUDA 13.4 (torch 2.14.1+cu130): BF16 inputs, head dim 128, chunk 64; forward + backward (f+b) and forward (fwd) per call, ms (median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR.
NVIDIA B200 (sm_100a) and NVIDIA GB300 (sm_103a): unchanged; the compiled programs are identical to FlashInfer main, the launch trace is identical, and outputs are bitwise equal. Deterministic, single-rank, sm_100a and sm_103aSource: flashinfer-ai/flashinfer#8615. Baseline: the det package on FlashInfer main before it. No non-Cake baseline. Its sm_100a and sm_103a programs are the vendored ones (see above).
NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.4 (torch 2.14.1+cu130): BF16 inputs, head dim 128, chunk 64; forward + backward (f+b) and forward (fwd) per call, ms (median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR.
NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4 (torch 2.14.1+cu130; driver 580.159.03): BF16 inputs, head dim 128, chunk 64; forward + backward (f+b) and forward (fwd) per call, ms (median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR.
Repeat measurement of the rows within noise above, on NVIDIA B200 (sm_100a, 148 SMs) and NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.4 (torch 2.14.1+cu130): fresh run on another node of the same GPU type, same settings (BF16 inputs, head dim 128, chunk 64; f+b per call, ms; median over 56 interleaved clean-process cells); speedup = FlashInfer main / this PR. Rows marked "within noise" have a 95 % confidence interval that includes 1.0.
|
Performance (2/5), continuedAs in part 1, "this PR", "new", "previous" and "main" in the copied FlashInfer text refer to that FlashInfer pull request and to the FlashInfer package before it, not to this PR. Deterministic KCPExact revision: flashinfer-ai/flashinfer#8585This is a correctness fix with no kernel or launch change. Its tables compare the fix with FlashInfer main before it. With 8 GPUs (2 nodes x 4), context parallel 8; four documents of 98,304 / 65,536 / 57,344 / 40,960 tokens, 32,768 tokens per rank; (H_qk, HV) = (16, 48); median ms of the slowest rank over interleaved main/fix visits; CUDA 13.4 toolkit, PyTorch 2.14.1+cu130. GB300 (sm_103a)
On the contiguous layout, P = 3 on every rank, so main's gradients are wrong in the GB200 (sm_100a)
Same kernels: flashinfer-ai/flashinfer#8538A = the det KCP package after flashinfer-ai/flashinfer#8385; B = flashinfer-ai/flashinfer#8538. B/A below 1 means B is faster. Cake's deterministic context-parallel GDN training kernel has one data-type variant: BF16 q/k/v/out/dout with FP32 scalar log-decay g and beta. The grouped Q/K head counts ms = slowest-rank median; B1 = module forward with autograd recording; B3 = forward + non-reentrant checkpoint replay + backward, including all communication and copies; Walk forms = changed (this PR launches a different width) or unchanged (same launches as A); s1/s2 = two sessions. Packs: P1 241,817 tokens / 6 documents, P2 261,987 / 3, P3 259,174 / 2, P4 254,976 / 1 (one document spans the whole group), P5 249,856 / 4. Heads (16, 48); causal conv1d width 4 over 10,240 channels. Software: CUDA 13.4 toolkit, torch 2.14.1+cu130. NVIDIA GB300 (sm_103a)
NVIDIA GB200 (sm_100a)
|
Performance (3/5), continuedAs in part 1, "this PR", "new", "previous" and "main" in the copied FlashInfer text refer to that FlashInfer pull request and to the FlashInfer package before it, not to this PR. Fast (non-deterministic), single-rank, sm_100a and sm_103a (exact revision)Source: flashinfer-ai/flashinfer#8675. Its "deterministic Cake" arm has the vendored det sm_100a and sm_103a programs. Cake ships four GDN training variants in FlashInfer: a non-deterministic and a deterministic form, each as a single-rank package and a context-parallel (KCP) package. This PR changes only the non-deterministic single-rank package Ratios are baseline ms / this PR ms, so a value above 1 means this PR is faster. ckpt is the checkpointed step. Before this PR, four cells were slower than deterministic Cake beyond run-to-run noise: B200 FP32-beta 6 seqs, 241817 tokens fwd+bwd; B200 FP32-beta 64 x 4096 fwd+bwd and ckpt; and GB300 FP32-beta 64 x 4096 ckpt. All four are now at parity with deterministic Cake (0.998x-1.004x). Compared with the previous non-deterministic package, they are 1.007x-1.019x. Two of them (B200 64 x 4096 fwd+bwd and GB300 64 x 4096 ckpt) are classed faster, and the other two are at parity. The bwd rows show where this PR costs time. On the four group-of-three rows, the isolated backward is 0.977x-0.991x of the previous non-deterministic package on both GPUs (0.9-2.4 % more time), which is the same time as the deterministic package (0.999x-1.002x). On those rows, fwd+bwd and ckpt are at parity with the previous package or faster. Two BF16-beta fwd+bwd cells are still slower than deterministic Cake beyond run-to-run noise, both at 7 x ~4.3k (real pack), H/HV 8/16: 0.985x on B200 and 0.992x on GB300. This PR does not change the schedule of those rows, and they are at parity with the previous package. The tables use three baselines:
All timings are eager, on NVIDIA B200 (sm_100a, 148 SMs) and NVIDIA GB300 (sm_103a, 152 SMs), with CUDA 13.0 and torch 2.14.1+cu130. NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.0 (torch 2.14.1+cu130): BF16 q / k / v, FP32 beta, head dim 128, chunk 64, eager; ms per call = mean of two runs in opposite arm orders, each the median of 6 group medians of 16 samples (L2 flushed before each sample); speedup = baseline ms / this PR ms; bwd = the backward pass alone (its forward runs outside the timed region), fwd+bwd = one training step, ckpt = forward, recomputed forward and backward of a checkpointed step.
NVIDIA B200 (sm_100a, 148 SMs), CUDA 13.0 (torch 2.14.1+cu130): BF16 q / k / v, BF16 beta, head dim 128, chunk 64, eager; ms per call = mean of two runs in opposite arm orders, each the median of 6 group medians of 16 samples (L2 flushed before each sample); speedup = baseline ms / this PR ms; bwd = the backward pass alone (its forward runs outside the timed region), fwd+bwd = one training step, ckpt = forward, recomputed forward and backward of a checkpointed step.
NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.0 (torch 2.14.1+cu130): BF16 q / k / v, FP32 beta, head dim 128, chunk 64, eager; ms per call = mean of two runs in opposite arm orders, each the median of 6 group medians of 16 samples (L2 flushed before each sample); speedup = baseline ms / this PR ms; bwd = the backward pass alone (its forward runs outside the timed region), fwd+bwd = one training step, ckpt = forward, recomputed forward and backward of a checkpointed step.
NVIDIA GB300 (sm_103a, 152 SMs), CUDA 13.0 (torch 2.14.1+cu130): BF16 q / k / v, BF16 beta, head dim 128, chunk 64, eager; ms per call = mean of two runs in opposite arm orders, each the median of 6 group medians of 16 samples (L2 flushed before each sample); speedup = baseline ms / this PR ms; bwd = the backward pass alone (its forward runs outside the timed region), fwd+bwd = one training step, ckpt = forward, recomputed forward and backward of a checkpointed step.
|
Performance (4/5), continuedAs in part 1, "this PR", "new", "previous" and "main" in the copied FlashInfer text refer to that FlashInfer pull request and to the FlashInfer package before it, not to this PR. Fast (non-deterministic), single-rank, sm_107aSource: the R200 table of flashinfer-ai/flashinfer#8665. Its "this PR" column runs the vendored sm_107a programs and schedule. Its "deterministic Cake" column is the det package at FlashInfer main Every ratio is baseline ms / this PR ms, so a value above 1 means this PR is faster. ckpt is the checkpointed step. The tables use three baselines:
R200 was timed with FP32 beta only; this PR does not change the R200 program. On R200 (sm_107a), the package keeps the previous WY program, and every cell is at parity with the previous version (0.994-1.001). Against deterministic Cake, it is 1.000-1.186. Against non-Cake SOTA, the ratio is above 1 in all 19 fwd+bwd cells (1.258-2.527) and all 14 published ckpt cells (1.106-2.066). In fwd, it is below 1 in 12 of 16 published cells (0.452-0.860) and above 1 in the other 4 (1.155-1.291). NVIDIA R200 (sm_107a, 212 SMs), CUDA 13.0 (torch 2.14.1+cu130): BF16 q / k / v, FP32 beta, head dim 128, chunk 64, eager; ms per call = mean of two runs in opposite arm orders, each the median of 6 group medians of 16 samples (L2 flushed before each sample); speedup = baseline ms / this PR ms.
Fast (non-deterministic) KCPThe default of this layer is flashinfer-ai/flashinfer#8644: eager,
|
| shape | documents (tokens) | tokens | heads (q/k, v) |
|---|---|---|---|
| A | 4: 98304/65536/57344/40960 | 262144 | (16, 48) |
| G | 12: 3598/4930/6262/7593/8925/10256/11588/12920/14251/15583/16914/18252 | 131072 | (16, 48) |
| H | 4: 49152/32768/28672/20480 | 131072 | (16, 48) |
| I | 4: 24576/16384/14336/10240 | 65536 | (16, 48) |
| M | 5: 12000/9000/6000/3768/2000 | 32768 | (16, 48) |
Legend: '*' = ratio below 1.000 (this PR slower than that baseline). Ratios = baseline ms / this PR ms from the shown absolutes. '(sustained)' = a cell whose short run hit the GPU software power cap (>= 5 % of samples), re-measured under the sustained protocol (idle cool-down before each block, 300-sample blocks, 3 block pairs per arm in both orders; value = median of the block medians); 'power-capped' = capped short-run cell without a sustained re-measure: no latency is claimed.
B. Two baselines (det(main), this PR and the in-process non-Cake arms in one run; ‡ = FlashQLA in its own process)
Strongest non-Cake = the fastest arm that passes the accuracy check: FLA-CP = flash-linear-attention 0.6.0's own context parallelism; FLA-HS = head split (all-to-all) around FLA's single-GPU kernels; FROST-ck0 / -ck64 = the same head split around cuDNN frontend 1.32.1's GDN (states recomputed / stored every 64 tokens). A non-Cake cell is published when this PR's arm and the shown arm drift <= 2 % between groups and every arm with drift > 2 % is slower than the shown arm in every group. At CP > 1 the strongest applicable non-Cake arm is FlashQLA 0.1.3 dual-CP (inter + intra rank context parallelism, contiguous layout): it runs in its own process and environment, paired with FLA-CP there and checked against it (relative RMS <= 1e-2, max abs <= 2e-2 of the reference max); '‡' = a FlashQLA value, its ratio divides by this PR's ms of the same table-B run, and the FLA-CP anchor of both processes is listed under the table. A CP > 1 cell without a FlashQLA value (layout D: FlashQLA has no per-document zigzag layout; a GPU without a FlashQLA build) is not claimed; the best in-process arm is shown as a note only.
| GPU / arch | workload (shape, CP, layout) | op | det(main) Cake ms | strongest non-Cake ms (arm) | this PR ms | det(main) / this PR | non-Cake / this PR |
|---|---|---|---|---|---|---|---|
| B200 sm_100a | A CP1 C | fwd (sustained) | 11.914 | 3.866 (FROST-ck0) | 11.104 | 1.073 | 0.348* |
| B200 sm_100a | A CP1 C | fwd+bwd (sustained) | 30.379 | 20.784 (FROST-ck64) | 29.592 | 1.027 | 0.702* |
| B200 sm_100a | A CP1 C | ckpt (sustained) | 39.761 | 25.885 (FROST-ck64) | 38.944 | 1.021 | 0.665* |
| B200 sm_100a | G CP1 C | fwd (sustained) | 5.980 | 1.861 (FROST-ck0) | 5.324 | 1.123 | 0.350* |
| B200 sm_100a | G CP1 C | fwd+bwd (sustained) | 15.295 | 10.642 (FROST-ck64) | 14.631 | 1.045 | 0.727* |
| B200 sm_100a | G CP1 C | ckpt (sustained) | 19.735 | 13.097 (FROST-ck64) | 19.437 | 1.015 | 0.674* |
| B200 sm_100a | M CP1 C | fwd | 1.695 | 0.594 (FROST-ck0) | 1.521 | 1.114 | 0.391* |
| B200 sm_100a | M CP1 C | fwd+bwd (sustained) | 4.077 | 3.244 (FROST-ck64) | 3.897 | 1.046 | 0.832* |
| B200 sm_100a | M CP1 C | ckpt | 5.411 | 4.238 (FROST-ck64) | 5.248 | 1.031 | 0.808* |
| B200 sm_100a | A CP8 C | fwd | 2.583 | 1.828 (FlashQLA 0.1.3 dual-CP‡) | 2.384 | 1.083 | 0.767*‡ |
| B200 sm_100a | A CP8 C | fwd+bwd | 5.322 | 5.749 (FlashQLA 0.1.3 dual-CP‡) | 5.117 | 1.040 | 1.124‡ |
| B200 sm_100a | A CP8 C | ckpt | 7.495 | 7.525 (FlashQLA 0.1.3 dual-CP‡) | 7.277 | 1.030 | 1.034‡ |
| B200 sm_100a | A CP8 D | fwd | 2.597 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 1.828 ms); best in-process arm 4.041 (FROST-ck0) | 2.424 | 1.071 | not claimed |
| B200 sm_100a | A CP8 D | fwd+bwd | 5.518 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 5.749 ms); best in-process arm 10.234 (FROST-ck64) | 5.335 | 1.034 | not claimed |
| B200 sm_100a | A CP8 D | ckpt | 7.715 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 7.525 ms); best in-process arm 13.407 (FROST-ck64) | 7.512 | 1.027 | not claimed |
| B200 sm_100a | H CP8 C | fwd | 1.542 | 1.142 (FlashQLA 0.1.3 dual-CP‡) | 1.552 | 0.994* | 0.736*‡ |
| B200 sm_100a | H CP8 C | fwd+bwd | 3.720 | 3.316 (FlashQLA 0.1.3 dual-CP‡) | 3.180 | 1.170 | 1.043‡ |
| B200 sm_100a | H CP8 C | ckpt | 5.421 | 4.645 (FlashQLA 0.1.3 dual-CP‡) | 4.475 | 1.211 | 1.038‡ |
| B200 sm_100a | H CP8 D | fwd | 1.604 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 1.142 ms); best in-process arm 2.192 (FROST-ck0) | 1.491 | 1.076 | not claimed |
| B200 sm_100a | H CP8 D | fwd+bwd | 3.489 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 3.316 ms); best in-process arm 5.457 (FROST-ck64) | 3.117 | 1.119 | not claimed |
| B200 sm_100a | H CP8 D | ckpt | 5.050 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 4.645 ms); best in-process arm 7.250 (FROST-ck64) | 4.428 | 1.140 | not claimed |
| B200 sm_100a | I CP8 C | fwd | 1.268 | 0.830 (FlashQLA 0.1.3 dual-CP‡) | 0.978 | 1.297 | 0.849*‡ |
| B200 sm_100a | I CP8 C | fwd+bwd | 3.118 | 2.426 (FlashQLA 0.1.3 dual-CP‡) | 2.262 | 1.378 | 1.073‡ |
| B200 sm_100a | I CP8 C | ckpt | 4.777 | 3.593 (FlashQLA 0.1.3 dual-CP‡) | 3.456 | 1.382 | 1.040‡ |
| B200 sm_100a | I CP8 D | fwd | 1.169 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 0.830 ms); best in-process arm 1.248 (FROST-ck0) | 1.030 | 1.135 | not claimed |
| B200 sm_100a | I CP8 D | fwd+bwd | 3.018 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 2.426 ms); best in-process arm 3.036 (FROST-ck64) | 2.233 | 1.352 | not claimed |
| B200 sm_100a | I CP8 D | ckpt | 4.585 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 3.593 ms); best in-process arm 4.186 (FROST-ck64) | 3.388 | 1.353 | not claimed |
| GB300 sm_103a | A CP1 C | fwd | 11.191 | 3.127 (FROST-ck0) | 10.285 | 1.088 | 0.304* |
| GB300 sm_103a | A CP1 C | fwd+bwd | 28.106 | 18.842 (FROST-ck64) | 27.086 | 1.038 | 0.696* |
| GB300 sm_103a | A CP1 C | ckpt | 36.914 | 23.159 (FROST-ck64) | 35.896 | 1.028 | 0.645* |
| GB300 sm_103a | G CP1 C | fwd | 5.659 | 1.655 (FROST-ck0) | 4.961 | 1.141 | 0.334* |
| GB300 sm_103a | G CP1 C | fwd+bwd | 13.868 | 9.817 (FROST-ck64) | 13.186 | 1.052 | 0.745* |
| GB300 sm_103a | G CP1 C | ckpt | 18.246 | 12.032 (FROST-ck64) | 17.569 | 1.039 | 0.685* |
| GB300 sm_103a | M CP1 C | fwd | 1.702 | 0.651 (FROST-ck0) | 1.502 | 1.133 | 0.433* |
| GB300 sm_103a | M CP1 C | fwd+bwd | 3.859 | 3.101 (FROST-ck64) | 3.662 | 1.054 | 0.847* |
| GB300 sm_103a | M CP1 C | ckpt (re-measure) | 5.126 | 4.169 (FROST-ck64) | 4.894 | 1.047 | not published (drift) |
| GB300 sm_103a | A CP4 C | fwd | 4.274 | 3.284 (FlashQLA 0.1.3 dual-CP‡) | 3.873 | 1.104 | 0.848*‡ |
| GB300 sm_103a | A CP4 C | fwd+bwd | 9.049 | 10.453 (FlashQLA 0.1.3 dual-CP‡) | 8.655 | 1.046 | 1.208‡ |
| GB300 sm_103a | A CP4 C | ckpt | 12.556 | 13.710 (FlashQLA 0.1.3 dual-CP‡) | 12.151 | 1.033 | 1.128‡ |
| GB300 sm_103a | A CP4 D | fwd | 4.331 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 3.284 ms); best in-process arm 7.253 (FROST-ck0) | 3.987 | 1.086 | not claimed |
| GB300 sm_103a | A CP4 D | fwd+bwd | 9.288 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 10.453 ms); best in-process arm 18.667 (FROST-ck64) | 8.957 | 1.037 | not claimed |
| GB300 sm_103a | A CP4 D | ckpt | 12.862 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 13.710 ms); best in-process arm 24.397 (FROST-ck64) | 12.525 | 1.027 | not claimed |
| GB300 sm_103a | H CP4 C | fwd | 2.712 | 1.883 (FlashQLA 0.1.3 dual-CP‡) | 2.552 | 1.063 | 0.738*‡ |
| GB300 sm_103a | H CP4 C | fwd+bwd | 5.450 | 5.584 (FlashQLA 0.1.3 dual-CP‡) | 5.150 | 1.058 | 1.084‡ |
| GB300 sm_103a | H CP4 C | ckpt | 7.772 | 7.395 (FlashQLA 0.1.3 dual-CP‡) | 7.333 | 1.060 | 1.008‡ |
| GB300 sm_103a | H CP4 D | fwd | 2.495 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 1.883 ms); best in-process arm 3.868 (FROST-ck0) | 2.276 | 1.096 | not claimed |
| GB300 sm_103a | H CP4 D | fwd+bwd | 5.110 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 5.584 ms); best in-process arm 9.775 (FROST-ck64) | 4.905 | 1.042 | not claimed |
| GB300 sm_103a | H CP4 D | ckpt | 7.140 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 7.395 ms); best in-process arm 12.852 (FROST-ck64) | 6.926 | 1.031 | not claimed |
| GB300 sm_103a | I CP4 C | fwd | 2.051 | 1.311 (FlashQLA 0.1.3 dual-CP‡) | 1.477 | 1.389 | 0.888*‡ |
| GB300 sm_103a | I CP4 C | fwd+bwd | 4.437 | 3.375 (FlashQLA 0.1.3 dual-CP‡) | 2.828 | 1.569 | 1.193‡ |
| GB300 sm_103a | I CP4 C | ckpt | 6.873 | 5.059 (FlashQLA 0.1.3 dual-CP‡) | 4.489 | 1.531 | 1.127‡ |
| GB300 sm_103a | I CP4 D | fwd | 1.827 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 1.311 ms); best in-process arm 2.085 (FROST-ck0) | 1.427 | 1.280 | not claimed |
| GB300 sm_103a | I CP4 D | fwd+bwd | 4.041 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 3.375 ms); best in-process arm 5.148 (FROST-ck64) | 2.889 | 1.399 | not claimed |
| GB300 sm_103a | I CP4 D | ckpt | 6.256 | not claimed (FlashQLA 0.1.3 dual-CP: no per-document zigzag layout; contiguous layout, same documents: 5.059 ms); best in-process arm 6.889 (FROST-ck64) | 4.393 | 1.424 | not claimed |
| R200 sm_107a | A CP1 C | fwd | 6.932 | 1.885 (FROST-ck0) | 6.663 | 1.040 | 0.283* |
| R200 sm_107a | A CP1 C | fwd+bwd | 17.688 | 16.279 (FROST-ck64) | 17.439 | 1.014 | 0.933* |
| R200 sm_107a | A CP1 C | ckpt | 23.345 | 18.930 (FROST-ck64) | 23.113 | 1.010 | 0.819* |
| R200 sm_107a | G CP1 C | fwd | 3.324 | 0.960 (FROST-ck0) | 2.948 | 1.128 | 0.326* |
| R200 sm_107a | G CP1 C | fwd+bwd | 8.642 | 8.415 (FROST-ck64) | 8.261 | 1.046 | 1.019 |
| R200 sm_107a | G CP1 C | ckpt | 11.316 | 9.732 (FROST-ck64) | 10.923 | 1.036 | 0.891* |
| R200 sm_107a | M CP1 C | fwd | 1.000 | 0.353 (FROST-ck0) | 0.896 | 1.116 | 0.394* |
| R200 sm_107a | M CP1 C | fwd+bwd | 2.381 | 2.334 (FROST-ck64) | 2.275 | 1.047 | 1.026 |
| R200 sm_107a | M CP1 C | ckpt | 3.167 | 2.847 (FROST-ck64) | 3.059 | 1.035 | 0.931* |
| R200 sm_107a | A CP4 C | fwd | 2.765 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 6.460 (FROST-ck0) | 2.560 | 1.080 | not claimed |
| R200 sm_107a | A CP4 C | fwd+bwd | 6.086 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 16.702 (FROST-ck64) | 5.885 | 1.034 | not claimed |
| R200 sm_107a | A CP4 C | ckpt | 8.375 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 21.441 (FROST-ck64) | 8.221 | 1.019 | not claimed |
| R200 sm_107a | A CP4 D | fwd | 2.780 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 6.460 (FROST-ck0) | 2.581 | 1.077 | not claimed |
| R200 sm_107a | A CP4 D | fwd+bwd | 5.957 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 16.707 (FROST-ck64) | 5.763 | 1.034 | not claimed |
| R200 sm_107a | A CP4 D | ckpt | 8.373 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 21.444 (FROST-ck64) | 8.170 | 1.025 | not claimed |
| R200 sm_107a | H CP4 C | fwd | 1.592 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 3.412 (FROST-ck0) | 1.422 | 1.120 | not claimed |
| R200 sm_107a | H CP4 C | fwd+bwd | 3.345 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 8.714 (FROST-ck64) | 3.042 | 1.100 | not claimed |
| R200 sm_107a | H CP4 C | ckpt | 4.723 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 11.248 (FROST-ck64) | 4.348 | 1.086 | not claimed |
| R200 sm_107a | H CP4 D | fwd | 1.578 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 3.414 (FROST-ck0) | 1.450 | 1.088 | not claimed |
| R200 sm_107a | H CP4 D | fwd+bwd | 3.297 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 8.714 (FROST-ck64) | 3.171 | 1.040 | not claimed |
| R200 sm_107a | H CP4 D | ckpt | 4.667 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 11.251 (FROST-ck64) | 4.533 | 1.030 | not claimed |
| R200 sm_107a | I CP4 C | fwd | 0.982 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 1.824 (FROST-ck0) | 0.864 | 1.137 | not claimed |
| R200 sm_107a | I CP4 C | fwd+bwd | 2.246 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 4.570 (FROST-ck64) | 1.776 | 1.265 | not claimed |
| R200 sm_107a | I CP4 C | ckpt | 3.566 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 5.965 (FROST-ck64) | 2.545 | 1.401 | not claimed |
| R200 sm_107a | I CP4 D | fwd | 0.980 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 1.822 (FROST-ck0) | 0.903 | 1.085 | not claimed |
| R200 sm_107a | I CP4 D | fwd+bwd | 1.989 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 4.566 (FROST-ck64) | 1.887 | 1.054 | not claimed |
| R200 sm_107a | I CP4 D | ckpt | 3.152 | not claimed (FlashQLA 0.1.3 dual-CP: not available for sm_107a); best in-process arm 5.964 (FROST-ck64) | 2.721 | 1.158 | not claimed |
FLA-CP anchor of the ‡ cells (FlashQLA process / table-B run, ms): B200 sm_100a A CP8 fwd 9.561 / 9.583; B200 sm_100a A CP8 fwd+bwd 25.743 / 25.808; B200 sm_100a A CP8 ckpt 35.250 / 35.297; B200 sm_100a H CP8 fwd 5.040 / 5.075; B200 sm_100a H CP8 fwd+bwd 13.189 / 13.209; B200 sm_100a H CP8 ckpt 18.071 / 18.099; B200 sm_100a I CP8 fwd 2.777 / 2.795; B200 sm_100a I CP8 fwd+bwd 7.009 / 7.013; B200 sm_100a I CP8 ckpt 9.603 / 9.615; GB300 sm_103a A CP4 fwd 17.655 / 18.049; GB300 sm_103a A CP4 fwd+bwd 41.095 / 41.809; GB300 sm_103a A CP4 ckpt 58.789 / 59.796; GB300 sm_103a H CP4 fwd 8.959 / 9.113; GB300 sm_103a H CP4 fwd+bwd 20.638 / 20.970; GB300 sm_103a H CP4 ckpt 29.454 / 29.939; GB300 sm_103a I CP4 fwd 4.768 / 4.855; GB300 sm_103a I CP4 fwd+bwd 10.629 / 10.784; GB300 sm_103a I CP4 ckpt 15.192 / 15.420.
B200 CP1 cells marked (sustained) in table B: one process per shape (A; G; M), each holding det(main), this PR, previous and the four in-process non-Cake arms, at the PR head.
Table A of flashinfer-ai/flashinfer#8644 (this PR against the previous fast KCP package, 81 cells, with a det(main) column from the same runs) is not copied here. Its summary:
"Previous" is this package at FlashInfer main f2926f20. Ratios are previous / this PR, so a value above 1 means this PR is faster. In table A, no cell is slower than previous on any GPU. On B200 and GB300, 23 of 27 cells are faster and 4 are at parity. On a 212-SM R200, all 27 cells are faster. The parity cells are all contiguous-layout cells of shapes H and I, at CP8 on B200 and CP4 on GB300. The forward gains the most: 1.004-1.104 on B200, 1.011-1.111 on GB300 and 1.024-1.103 on R200. On every GPU the best cell is shape G at CP1: 5.899 -> 5.344 ms on B200, 5.522 -> 4.972 ms on GB300 and 3.252 -> 2.947 ms on R200. fwd+bwd gains 1.003-1.042 on B200, 1.002-1.043 on GB300 and 1.011-1.038 on R200. The checkpointed step (ckpt) gains 1.004-1.041, 1.001-1.032 and 1.007-1.030 on the same three GPUs.
Performance (5/5), continued (fast KCP, continued; det kernels with FLA's context parallelism; missing measurements)As in part 1, "this PR", "new", "previous" and "main" in the copied FlashInfer text refer to that FlashInfer pull request and to the FlashInfer package before it, not to this PR. flashinfer-ai/flashinfer#8638 (exact revision): PR head, graph modedet(main) is the deterministic KCP package of the same main. Against it, this PR is faster or at parity in every published A4 cell:
Table B also compares this PR with the strongest accuracy-qualified non-Cake arm. Its ratios are non-Cake ms / this PR ms, so above 1 means this PR is faster. At CP1, the cuDNN frontend GDN arms (FROST) are faster than this PR on shapes G, A and M. There non-Cake / this PR is 0.396-0.866 on B200, 0.341-0.954 on GB300 and 0.316-0.938 on R200. The one exception is R200 shape M ckpt, where the ratio is 1.003. On shape N, this PR is faster where a cell is published: 1.888-1.902 on B200 and 1.562 on R200. At CP > 1, the strongest accuracy-qualified non-Cake arm is FlashQLA 0.1.3 with its own dual context parallelism. FlashQLA needs its own environment, so it runs in a separate process. FLA-CP runs in both processes and is the anchor between them. On B200 at CP8, FlashQLA / this PR is 0.887 in fwd, 0.928 in fwd+bwd and 1.102 in ckpt. On GB300 at CP4, it is 0.951, 0.983 and 1.208. FlashQLA is faster in fwd and fwd+bwd, and this PR is faster in ckpt. Table B has no CP > 1 rows on R200. Shapes (head dim 128, BF16 q / k / v, FP32 g / beta, in-kernel q / k L2 norm, layout contiguous, GDN layer only): A = 262144 tokens, 4 documents (98304 / 65536 / 57344 / 40960), heads (16, 48) (q/k heads, v heads); G = 131072 tokens, 12 documents (3598 ... 18252), heads (16, 48); M = 32768 tokens, 5 documents (12000 / 9000 / 6000 / 3768 / 2000), heads (16, 48); N = 8192 tokens, 3 documents (2049 / 3001 / 3142), heads (2, 6). CPn = context parallel over n GPUs. The single-rank proxy forces two segments per fragment (the segment walks and carry merges a CP > 1 fragment runs) on one GPU. Non-Cake arms: FLA-CP = flash-linear-attention 0.6.0's own context parallelism; FLA-HS = head split (all-to-all) around FLA's single-GPU kernels; FROST-ck0 / -ck64 = the same head split around cuDNN frontend 1.32.1's GDN (states recomputed / stored every 64 tokens). Marks used in these tables, from the legend of flashinfer-ai/flashinfer#8638: B200 (sm_100a): B. Two baselines at the PR head (graph mode for the Cake arms; the non-Cake arms have no graph form here and run eager; det(main) = the deterministic KCP package of main f2926f20)
FLA-CP anchor of the FlashQLA cells (FlashQLA process / two-baseline process, ms): A CP8 fwd 9.575 / 9.636; A CP8 fwd+bwd (sustained) 25.765 / 25.736; A CP8 ckpt (sustained) 35.246 / 35.191. A4. PR head (on FlashInfer main f2926f20, which holds flashinfer-ai/flashinfer#8600 and flashinfer-ai/flashinfer#8623): graph N / M fwd+bwd and checkpointed step (this PR vs previous = main f2926f20, same process, A/A arm, both arm orders twice)
Same runs, det(main) (the deterministic KCP package of main f2926f20) vs this PR:
GB300 (sm_103a): B. Two baselines at the PR head (graph mode for the Cake arms; the non-Cake arms have no graph form here and run eager; det(main) = the deterministic KCP package of main f2926f20)
FLA-CP anchor of the FlashQLA cells (FlashQLA process / two-baseline process, ms): A CP4 fwd 17.784 / 17.875; A CP4 fwd+bwd 41.132 / 41.349; A CP4 ckpt 59.121 / 59.003. A4. PR head (on FlashInfer main f2926f20, which holds flashinfer-ai/flashinfer#8600 and flashinfer-ai/flashinfer#8623): graph N / M fwd+bwd and checkpointed step (this PR vs previous = main f2926f20, same process, A/A arm, both arm orders twice)
Same runs, det(main) (the deterministic KCP package of main f2926f20) vs this PR:
R200 (sm_107a, 212 SMs): B. Two baselines at the PR head (graph mode for the Cake arms; the non-Cake arms have no graph form here and run eager; det(main) = the deterministic KCP package of main f2926f20), on a 212-SM board
A4. PR head on a 212-SM board (on FlashInfer main f2926f20, which holds flashinfer-ai/flashinfer#8600 and flashinfer-ai/flashinfer#8623): graph N / M fwd+bwd and checkpointed step (this PR vs previous = main f2926f20, same process, A/A arm, both arm orders twice)
Same runs, det(main) (the deterministic KCP package of main f2926f20) vs this PR:
Tables B and A4 above time the PR head. Tables A, A2 and A3 of flashinfer-ai/flashinfer#8638 time earlier revisions of that PR and are not copied here. Deterministic kernels with FLA's context parallelism (radixark/miles#3597)These are the deterministic chunked GDN training kernels of radixark/miles#3597 (open; head Revision note: the PR body's whole-model runs and kernel gates are at branch head Setup from the PR body: a 27B dense GDN model, GRPO, Sampled arms. Means over steps >= 1. Times are in seconds.
Three rollouts were recorded with cake. They were then replayed train-only twice with cake (a, b) and once with fla (f). The cake value is the mean of a and b. Means are over all 3 steps, in seconds;
With step 0 excluded (steps 1–2), the cake/fla ratios on B200 are 0.975 for The table below lists kernel-level gains per round. Each gain is the CUPTI GPU time of the GDN chunk kernels (forward, and the forward+backward union) on the best row of the 16-row kernel table, from a same-node interleaved A/B against the previous round's kernels. Round 12 kept no candidate.
Missing facts
|
The README repeated the docstrings of _runtime/config.py, workspace.py, loader.py, registry.prebuild and tools/vendor_cake.py, which would drift from the code. Keep the install, miles contract, layout and adding-an-op sections, shorten the rest to pointers into those modules, and keep the Ops and Prebuild tables (test_op_contract reads the Ops table; later layers append rows) with their surrounding lines unchanged so the layers merge without conflicts. The version bump is not needed for this PR.
Install, the miles contract in one paragraph, the Ops and Prebuild tables (test_op_contract reads the Ops table; later layers append rows), the layout and the development commands. Everything else is documented in the runtime modules and tools/vendor_cake.py.
workspace.take read the current device, stream and capture state through private torch._C functions and kept several identity-keyed fast paths (_direct, _bound, _devices, a running byte total) to shave host time. The pool is off by default and measured no faster with it on, so the extra state is not worth its cost: read torch.cuda's public API, compute held bytes from the pool, and keep the same take / release / stats behaviour (one buffer per op, role and stream; LRU eviction under the bound; capture bypass). workspace.py goes from 248 to 130 lines. config.variant / config.option no longer memoise per configuration; a dict lookup and a scan over an op's few flags is enough. Tests: drop the cases that pinned the removed internals and the micro-benchmark bench_workspace.py; the fake CUDA layer patches the public torch.cuda functions only. On GB300 (torch 2.14.1+cu130) the runtime suite passes: 214 passed, 1 skipped (was 223, minus the 9 internal-only cases). CI: the upstream vendoring check fetches every recorded commit, so it moves to its own workflow that runs when VENDORED.json or the tool changes, weekly, and on demand.
main's tree equals cake-stack/00-layout at d67ad25, which this branch already contains, so this records main as merged without changing files.
…k/01-gdn # Conflicts: # README.md
Stack 1/13 · base: main
Summary
Adds the
gdnop,miles_ext_kernels.cake.gdn: the chunked gated delta rule of GDN linear-attention layers, forward and backward, on SM100a, SM103a and SM107a.chunk_gated_delta_rulekeeps FLA v0.5.2's signature and autograd contract.kcp_chunk_gated_delta_ruleandbuild_kcp_contextexpose the packages' own context parallelism (KCP).The two variants,
detandfast, run on five vendored packages:eb3eb68d, byte for byte: det and fast, each without context parallelism and with KCP.cake_gdn_fla_cp_det: the kernels of cake_gdn: deterministic chunked GDN training kernels (--linear-attention-backend cake) for the head-sharded Qwen GDN layer miles#3597, packaged as the det version of FLA's context-parallel (FLA-CP) path and built by the runtime loader.miles selects the op with
--use-miles-ext-kernels --linear-attention-backend cakeand importsfrom miles_ext_kernels.cake.gdn import chunk_gated_delta_rule(paired PR radixark/miles#4676, whose update to this layout follows this stack). Through miles' arguments, an unset--cake-gdn-variantresolves tofast, or todetunder--deterministic-mode. An explicitfastunder deterministic mode is rejected. A process that never callsconfigurerunsdet. Package scratch goes through the runtime pool, which keeps nothing at its default bound of 0.Stack
Layer 1 of 13. This PR replaces its previous version, "[Cake 1/11] gdn: shared build/vendor infra and GDN kernels (det, fast)" (head
c62be504). The shared infrastructure of that version moved to layer 0 ([Cake 0/13] layout); this layer is the gdn op only.Builds on layer 0 (head
e0482323). Every PR targets main, so until layer 0 merges, this diff also contains layer 0's 3 commits. This layer's commits (e0482323..ad8eeef8):e9dd400evendor: gdn packages631ab4b2gdn: chunked gated delta rule on the cake runtimead8eeef8tests: gdnChanges
This layer: 373 files, +316554 / -2 (previous version of this PR: 162 files, +106188 / -4). Per commit:
e9dd400e353 files, +309414 / -1;631ab4b28 files, +3020 / -1;ad8eeef812 files, +4120. Almost all of it is vendored:cake/_vendor/:cake_gdn_det_train85 (+67823),cake_gdn_det_train_kcp59 (+64357),cake_gdn_train80 (+67256),cake_gdn_train_kcp67 (+72224),cake_gdn_fla_cp_det61 = 57 upstream + 4 local (+37316)cake/_vendor/VENDORED.json(ops.gdnblock)cake/gdn.py65,cake/_impl/gdn/__init__.py25,ops.py2535,kcp.py229,native.py161;registry.py(OPS = ("gdn",), +1 / -1)tests/cake/gdn/pyproject.toml(gdnextra)The 56 kernel sources and
registry.jsonofcake_gdn_fla_cp_detkeep the blobs of the previous version'sgdn/csrcandgdn/registry.json. The previous version vendored onlycake_gdn_train, from FlashInferc5db8d3; the four FlashInfer packages here are ateb3eb68d. The message ofe9dd400elists both vendoring commands (with../flashinferand../milescheckouts).Packages
cake_gdn_det_traineb3eb68d7deebbe619e05d93ecfd9796e4eaf2e3,flashinfer/experimental/cake_gdn_det_traincake_gdn_trainflashinfer/experimental/cake/cake_gdn_traincake_gdn_det_train_kcpflashinfer/experimental/cake_gdn_det_train_kcpcake_gdn_train_kcpflashinfer/experimental/cake_gdn_train_kcpcake_gdn_fla_cp_det9d7543812b676a5b50dde107518782933b804b48,miles/kernels/attention/gdn/cake(csrc+registry.json)cp_context, and KCP calls the KCP packages do not serve), either variantmin_cudais 12.8 / 12.9 / 13.4, the same floor the runtime enforces for every build.eb3eb68dis on FlashInfer main and contains the context-parallel package fixes fix(cake_gdn_det_kcp): keep M_seg in the full saved set at P > 1 and count the whole plane in the saved byte query flashinfer-ai/flashinfer#8585 and fix(cake_gdn_kcp): byte queries of the non-deterministic KCP GDN package match what the calls bind flashinfer-ai/flashinfer#8623.cake_gdn_fla_cp_det: the 56 kernel/launcher sources andregistry.jsonare byte-identical to cake_gdn: deterministic chunked GDN training kernels (--linear-attention-backend cake) for the head-sharded Qwen GDN layer miles#3597's head. Its upstreamcsrc/cake_host_shim.his excluded (blobfdc706d6720ced076315d4a11701f03798c96186).vendor_cake.py miles-pradds an empty__init__.py, the FlashInfer loadercake_build.py, the TVM-FFI host shim fromtools/templates/cake_host_shim.hand a generatedbuild.json. That file builds sm_107a from the sm_103a records with-std=c++17 -O3 --compiler-options=-fPIC --expt-relaxed-constexprand links with-shared -lcudart -lcuda. The host code is_impl/gdn/ops.py+native.py; cake_gdn: deterministic chunked GDN training kernels (--linear-attention-backend cake) for the head-sharded Qwen GDN layer miles#3597's_jit.py/ops.py/__init__.pyare not vendored.allowentries: no vendored module reads the environment or imports flashinfer.Variant routing
The variant is read when the call runs, and a backward runs its forward's choice.
cp_contextcake_gdn_det_traincake_gdn_trainKcpContextcake_gdn_det_train_kcpcake_gdn_train_kcpcp_context, or aKcpContextcall the KCP packages do not servecake_gdn_fla_cp_detcake_gdn_fla_cp_det(the det version of this path), logged once per reasonKcpContextcall goes tocake_gdn_fla_cp_detwhen it uses the in-kernel L2 norm without its records in the package build, or a partition beyond the package's plan caps. That route needs a contiguous partition whoseworld_sizedivides the token count; otherwise the call raises.warn_runs_det(the fast variant does not serve an FLA cp_context; running det).det. Through miles' arguments: an unset flag givesfast(written back byvalidate_args, delivered byconfigure(args.miles_ext_kernels_args)),--deterministic-modegivesdet, and an explicitfastunder deterministic mode raisesValueError.cp_contextthe contract is FLA's:initial_stateandoutput_final_stateraiseValueErrorwhen the call is made, and the context'scu_seqlens/cu_seqlens_cpureplace the arguments.--cake-gdn-recompute-state. Across those, results agree only within the accuracy budget. fast makes no promise about its bits; the tests compare it within the accuracy budgets only. The recomputed det backward is compared within the budgets of the saved one.API and flags
cake/gdn.pyis torch-free:OP = OpSpec(name="gdn", variants=("det", "fast"), selected_by=(("linear_attention_backend", "cake"),), packages=(the 5 above), exports=("KcpContext", "build_kcp_context", "chunk_gated_delta_rule", "kcp_chunk_gated_delta_rule"), flags=(--cake-gdn-recompute-state,)). Entry points load lazily fromcake._impl.gdn.chunk_gated_delta_rule(q, k, v, g, beta, scale=None, initial_state=None, output_final_state=False, cu_seqlens=None, cu_seqlens_cpu=None, use_qk_l2norm_in_kernel=False, cp_context=None, chunk_size=64, **unsupported): q/k[B, T, H, K], v[B, T, HV, V], K = V = 128,HV % H == 0, BF16 q/k/v, packed varlen (B = 1), FP32 initial/final state[N, HV, K, V], in-kernel q/k L2 norm. Rejected explicitly:use_gate_in_kernel, beta sigmoid fusion,state_v_first,head_first, chunk sizes other than 64.cu_seqlens_cpu(tuple of ints or CPU tensor) selects the host-table route, with no device-to-host sync. A call without tokens returns an empty output and passes the state through.kcp_chunk_gated_delta_rule(q, k, v, g, beta, kcp_context, scale=None, use_qk_l2norm_in_kernel=False, chunk_size=64)runs only the KCP package of the configured variant; a call it does not serve raisesNotImplementedError.build_kcp_context(cu_seqlens_global, group=None, *, device=None, conv_kernel_size=4, layout="contiguous", fragments=None) -> KcpContextbuilds host metadata only (no tensor read, no collective). Layouts:contiguous(FLA's partition whenworld_sizedivides the token count),zigzag,zigzag_per_documentor explicitfragments. LRU cache of 64 contexts. KCP takes FP32 g/beta (BF16 beta is cast exactly) and no initial or final state.--cake-gdn-recompute-state(store_true, default off) recomputes w / u / h / v_new in the backward. The single-rank packages take it assave_state, the KCP packages assave_mode="light"; its help states the accuracy and the bits per route. This flag and--cake-gdn-variantare both rejected without--linear-attention-backend cake.Changes from the previous version of this PR
fastunder--deterministic-mode, or throughconfigure(..., deterministic_mode=True), is now rejected: fast does not promise its bits, so deterministic mode runs det. Before, it was accepted.cp_context. Before, a call with acp_contextused the caller'scu_seqlens(the global boundaries gave non-finite outputs) and acceptedoutput_final_state, returning the local final state with a backward that raised.cake_gdn_det_train._jit.py(torchcpp_extension), the module-levelVARIANT/SAVE_STATE/WORKSPACES, and the family's ownarguments.py/configure.cake_gdn_fla_cp_detgo through the runtime's workspace pool instead of per-package pools kept until the process ends. At the default--cake-workspace-pool-bytesof 0, each is a fresh allocation per call. Under a positive bound, the pool keeps one buffer per role and stream, freed byrelease(). The 0-element state placeholder is one tensor per device, shape and dtype; no call allocates one.gdn = ["flash-linear-attention>=0.5.2"]. The FLA-CP path calls thefla.ops.cphelpers of that release (with theirstate_v_firstkeyword) and is checked against it. 0.5.0's helpers lack the keyword, and 0.5.1's helpers module differs from 0.5.2's.Prebuild
cake.prebuild("gdn", *, device=None, max_workers=None)takes no model parameter. It builds every library of the five packages for the device's architecture (138 / 138 / 137 for sm_100a / sm_103a / sm_107a). It raisesNotImplementedErroroff SM100a / SM103a / SM107a, andRuntimeErrorwhen nvcc is too old (sm_107a needs CUDA 13.4). The memo key includescake_gdn_recompute_state.One exception, stated in the README Prebuild row and the docstring: FLA's own Triton context-parallel pre-processing (
fla.ops.cp) runs only when the caller passes an FLAcp_contextwith a process group. It is third-party code and compiles in FLA on its first call.Workspace and memory
workspace.take("gdn", role, ...): 14 roles inops.py(forwardf_seg_state,f_seg_flag; backwardb_dv_local,b_dh,b_dv2,b_seg_state,b_seg_flag,b_dh_end,b_dq_hv,b_dk_hv,b_dw,b_dg1,b_dk2,b_dg2). The FlashInfer det / fast single-rank packages' scratch hook (cake_chain._POOLED) points at the same pool. At the default bound of 0, every role is a fresh allocation per call and nothing is kept between calls.f_zero_state/b_zero_state, for a launch without initial or final state) are one tensor per device, shape and dtype for the process, not pool entries. They hold no memory, so streams share them without a race;release()drops them.release(). The packages' own small per-layout chunk tables and plans stay.release()drops the FLA-CP chunk tables, plans and loaded launchers, the cachedKcpContextobjects, the state placeholders and the single-rank packages' self-zeroing buffers.release()returns 0 (B200, GB300). With 1 GiB, det keeps 0.232 GiB in 5 entries (B200; GB300 249037952 B, 5 entries), det with FLA context parallelism 0.512 GiB in 11 entries (B200) and fast 0.232 GiB (B200);release()returns exactly the bytes held. det outputs and gradients are bit-identical to bound 0 (B200: det and det with FLA context parallelism; GB300: det, 12/12 tensors over 2 layers, on the miss and the hit run).Review items
From the review of this PR's previous version. The layout comment and the stack-wide mechanisms are resolved in layer 0.
cake_gdn_det_train(its own caps); the FLA-CP planner caps segments atPRESCAN_MAX_CHUNKS = 128and keeps the serial scan when it cannot split within ittest_gdn_host.py::test_one_pre_scan_cta_covers_at_most_128_chunks,::test_without_the_cap_a_long_single_sequence_would_cut_segments_the_pre_scan_cannot_serve,::test_no_plan_cuts_a_segment_longer_than_the_cap,::test_a_lower_capacity_bounds_the_split; GPUtest_gdn_det_gpu.py::test_a_long_single_sequence_runs_on_the_det_package_and_the_fla_cp_det_kernels(48v_32k, 12v_64k, 16v_128k; host and device),test_gdn_wiring.py::test_a_long_single_sequence_never_plans_a_segment_beyond_the_pre_scan,::test_segment_planner_only_splits_for_a_modelled_gain; ROUTES*_long_*_check_inputs(metadata only, no sync), and host boundaries must end at q's tokens, checked before any CUDA calltest_gdn_host.py::test_host_boundaries_must_end_at_the_tokens_before_any_cuda_call(tuple, tensor); GPUtest_gdn_det_gpu.py::test_the_extents_are_checked_before_a_launch(with and without CP; the well-posed follow-up call uses in-kernel q/k L2 norm, as miles calls it)cp_context, initial/final state are rejected at the call and the context's boundaries replace the arguments; the autograd function no longer materialises a missing incoming gradient (set_materialize_grads(False))test_gdn_kcp_wiring.py::test_the_context_parallel_contract_is_fla_s,::test_an_fla_context_runs_the_fla_cp_det_kernels_with_the_context_s_boundaries; GPUtest_gdn_fla_cp_gpu.py::test_the_context_s_boundaries_replace_the_call_s,::test_an_unused_final_state_passes_no_gradient_and_a_final_state_alone_backpropagates,test_gdn_kcp_gpu.py::test_an_initial_or_a_final_state_is_rejected_when_the_call_is_made;multi_gpu::test_cp_a_final_state_call_is_rejected_when_the_call_is_made(cp2, cp4)release(); out-of-pool tail buffers documented; 0-element placeholders cached instead of allocatedtest_gdn_wiring.py::test_at_the_default_bound_of_0_a_training_step_retains_no_pooled_bytes,::test_workspace_pool_hits_after_the_first_call_and_results_match_the_pool_off,::test_pooled_workspaces_never_alias_saved_returned_or_gradient_tensors,::test_a_side_stream_gets_its_own_buffers_and_the_same_bits,test_gdn_det_gpu.py::test_the_scratch_comes_from_the_bounded_pool_and_the_pool_does_not_change_the_bits,test_prebuild_gpu.py::test_a_det_call_keeps_its_scratch_in_the_bounded_pool(pool tests configurehelpers.POOLED); CPUtest_gdn_host.py::test_the_state_placeholder_is_one_tensor_per_device_shape_and_dtype,::test_the_packages_take_their_scratch_from_the_pool_and_share_the_state_placeholders(det, fast)gdn.py,ops.pyand the README (see Variant routing)test_arguments.py::test_the_recompute_help_states_the_accuracy_and_the_bitstest_gdn_host.py::test_a_call_without_tokens_passes_the_state_through,::test_a_call_without_tokens_and_without_a_final_state,::test_a_call_without_tokens_checks_the_initial_state; GPUtest_gdn_det_gpu.py::test_a_call_without_tokens(det and fast, with and without CP); ROUTESplain_no_token,fla_cp_no_tokengdn/csrc(58 files) outside every manifest, a second build system (cpp_extension)cake_gdn_fla_cp_det, drift-checked by--check --upstream;_jit.py/cpp_extensionremoved,native.pyloads loader-built tvm-ffi librariestest_package.py::test_the_fla_cp_package_builds_one_library_per_stage_record_and_architecture,::test_the_fla_cp_sources_build_with_the_flags_of_the_previous_build,::test_every_stage_has_the_sm103a_export_sm107a_compiles; GPUtest_gdn_fla_cp_gpu.py::test_every_stage_loads_from_the_loader_s_libraryVARIANT,SAVE_STATE,WORKSPACESprocess-global, set byconfiguretest_arguments.py::test_configure_sets_the_variant_and_the_recompute_switch_the_kernels_read; GPUtest_prebuild_gpu.py::test_a_call_runs_the_variant_configured_when_it_is_madepyproject.tomldeclares no dependenciesgdn = ["flash-linear-attention>=0.5.2"]prebuildbuilds all five packages, flags go throughOPtest_arguments.py::test_without_the_arguments_the_kernels_keep_the_tape_and_run_det,::test_an_explicit_fast_is_kept_and_rejected_under_deterministic_mode; GPUtest_prebuild_gpu.py::test_after_one_prebuild_no_route_compiles_vendor, one runtimecake/_vendor/and itsops.gdnblocktest_package.py::test_the_implementation_loads_exactly_the_declared_packagesAlso fixed after GPU runs of earlier versions (not review items): with the shared multi-GPU fixture, a case that raises on every rank before any communication fails only its own tests. The process group has a 5-minute timeout, so a case that raises on some ranks only cannot hang forever, but it is not isolated. Test-only fixes from the GPU rounds of this stack: the extents test's well-posed call normalises q/k in the kernel (unnormalised unit-variance inputs make the state diverge; FLA's Triton kernel also gives non-finite outputs there); planner tests plan at the B200 / GB300 SM counts their expectations were modelled for; the planner-gain test only considers lengths within the 128-chunk cap; the launch-count test skips, naming the device, where torch.profiler records no CUDA API call.
Tests
tests/cake/gdn/has 119 test functions. CPU:test_arguments.py13,test_package.py6,test_gdn_host.py10,test_gdn_kcp_wiring.py15. GPU (skipped without CUDA):test_gdn_wiring.py34,test_gdn_det_gpu.py12,test_gdn_fast_gpu.py5,test_gdn_fla_cp_gpu.py7,test_gdn_kcp_gpu.py12 (includingmulti_gpuCP 2 / 4),test_fla_parity.py1,test_prebuild_gpu.py4. On a GPU node, 318 are collected withoutmulti_gpuand 30 with it.helpers.POOLEDat the bits of bound 0, and bound 0 keeps no pooled byte after a training step; FLA-CP from the loader's libraries; KCP on one GPU andmulti_gpuCP 2 / 4.test_prebuild_gpu.pydeclares 23ROUTES:plain_host,plain_device,plain_one_sequence,plain_long_host,plain_long_device,plain_state,plain_state_device,plain_recompute,plain_no_l2norm_fp32_beta,plain_batch,plain_no_token,kcp,kcp_long,kcp_recompute,kcp_no_l2norm_fp32_beta,fla_cp_host,fla_cp_device,fla_cp_long_host,fla_cp_long_device,fla_cp_recompute,fla_cp_no_l2norm_fp32_beta,fla_cp_no_token,side_stream. Layout 300 / 517 / 64 / 7 tokens; long = one 9000-token sequence (141 chunks); 4 key and 12 value heads.check_prebuildruns every route under each variant with no build allowed. Planner expectations use the 148 / 152 SM counts of B200 / GB300.Validation
CPU: image
radixark/miles:dev-202611032236(Python 3.12.16, torch 2.14.1+cu130, tvm-ffi 0.1.11), x86_64 CPU node, CUDA hidden. The CPU tests compile nothing.pytest tests, fresh venv without torchpytest tests, image Pythonvendor_cake.py --checkpre-commit run --all-filese9dd400e(vendor)test_vendored.py::test_every_vendored_op_is_registered)631ab4b2(op)test_op_contract.py::test_the_op_has_cpu_tests_and_a_prebuild_test[gdn])ad8eeef8(tests)--check gdn --upstreamrc 0tests/cake/gdnalone in the image atad8eeef8: 158 passed, 7 skipped.--check --upstreamalso returns 0 at the stack top, whoseops.gdnblock,gdn.py,_impl/gdn/,tests/cake/gdn/and five package trees are identical to this layer's. GitHub CI has not run on this head yet.GPU:
tests/cake/gdnat the stack top (gdn files identical to this layer's), with a private JIT cache per op.multi_gputest_prebuild_gpu.pyfirst)test_prebuild_gpu.pyalone, first, on an empty cacheno_build())multi_gpu(4 GPUs)test_cp_a_final_state_call_is_rejected_when_the_call_is_made(cp2, cp4)test_gdn_det_gpu.py::test_det_repeats_bit_for_bit_in_a_new_process,test_gdn_fla_cp_gpu.py::test_the_call_repeats_bit_for_bit_in_a_new_processtest_gdn_det_gpu.py::test_det_repeats_on_a_side_stream_alternating_with_the_default_stream(both cases),test_gdn_fla_cp_gpu.py::test_a_side_stream_alternating_with_the_default_stream_gives_the_same_bits,test_gdn_wiring.py::test_a_side_stream_gets_its_own_buffers_and_the_same_bitssm_107a (R200): validation of this tree is pending.
Build migration (a one-time check, not a resident test; run on an earlier commit with the identical
cake_gdn_fla_cp_dettree): with nvcc 13.0.88, the loader-built FLA-CP libraries have the same SASS, function by function, as the previous torch-extension build for 50/50 libraries (25 per arch, sm_100a and sm_103a). With nvcc 13.4.92, sm_107a and sm_103a match 25/25 each against the previous compile flags. No torch-extension build for CUDA 13.4 was available, so the old side used its flags without torch. The whole-filecuobjdumptext differs only by one function-less ELF from the tvm-ffi binding unit. The old build's torch-only flags (-D/-U__CUDA_NO_HALF*,-std=c++20) do not change SASS.Performance
GDN det A/B: previous layout vs this stack (B200)
Same GPU, node otherwise idle. "old" is the previous layout of this op (per-op runtime, FlashInfer packages at
c5db8d3, samecake_gdn_fla_cp_detkernel sources); "new0" is this stack at the default pool bound 0; "new1g" is this stack with a 1 GiB pool. Arms: a = det kernels with FLA's context parallelism called as one rank (same kernel sources on both sides), b = det (FlashInfer det package), c = fast (FlashInfer fast package). Arms b and c also change the FlashInfer commit (c5db8d3->eb3eb68d), so their ratios include package changes, not only the layout.Method: eager, CUDA events, L2 flushed and device synchronized before every rep (host cost inside the window), 5 warm-up + 30 timed reps per process, arms interleaved ABC / CBA, 4 processes per side in rotated order; cell = median of the per-process medians; ratio [min, max] over the 4 rounds. Ops: fwd (with tape), bwd (after an untimed forward), fwdbwd, ckpt (forward under no_grad, forward with tape, backward). 15 shapes: the miles micro-batch layout (3339, 2802, 1340, 263 tokens) with 4x12 and 16x32 heads on host tables and 16x16 on device tables, plus 12 equal-length packed shapes (1/2/4 x 4096/8192 tokens, 8x16 and 16x32 heads). Torch 2.14.1+cu130, CUDA 13.0, driver 580.82.07. In all 12 processes the probes held: b and c bit for bit run to run, b
save_stateon/off same bits, and a / b / c within budget of an FP64 recurrence.Summary (median over rows; brackets = min and max over rows; ratios > 1 = this stack slower):
new0 / old by op (median): fwd 1.004, bwd 1.040, fwdbwd 1.038, ckpt 1.043. On the 4x8192-token shapes new0 / old is 0.971-1.046: arms b / c are faster in fwd, fwdbwd and ckpt (0.971-0.999) and slower in bwd (1.014-1.046). A 1 GiB pool is slower than bound 0, which supports the default of 0.
Selected rows (ms; 60 of the 180 rows):
The last column is an earlier run on another B200 node with a previous version of layer 0's
take(median new/old 1.039 there, miles shapes 1.064). A [max] near 2 comes from bimodal per-process bwd medians in single processes of the old or new1g side (e.g. old 0.22 vs 0.56 ms); the cells use the median over the 4 rounds.Host time per call under cProfile (miles-4x12-host, us, fwd / bwd; launch-return time, profiler overhead included):
workspace.takecumulative per op at bound 0: a bwd 12 calls 72.4 us, b bwd 6 calls 45.6 us (1 GiB: 130.6 / 78.0 us). Residual: a small host overhead remains on host-bound small shapes (miles shapes median 1.051x, max 1.185x on one equal-length ckpt row); on the 4x8192-token shapes arm a is within 1.013x, and arms b / c are faster except in bwd (1.014-1.046x).Vendored kernels (follow-up comments)
The performance of the vendored packages is copied from their source pull requests, with tables and captions unchanged, into five follow-up comments, "Performance (1/5)" to "Performance (5/5)":
save_mode="full", this layer's default (perf(cake_gdn_kcp): with grad, full-save forward writes the output and the tapes in one walk flashinfer-ai/flashinfer#8644).Missing, as listed in comment 5: a det vs non-Cake ratio (single-rank and KCP), det KCP on sm_107a against an earlier version, fast single-rank on sm_107a with BF16 beta, fast KCP at CP > 1 on R200 against a non-Cake implementation, fast KCP with
save_mode="light"in eager mode, for the FLA-CP det kernels any sm_107a or kernel-level non-Cake measurement and any measurement of their later WY backward change, and B300 numbers.Notes
cake_gdn_trainlists 34 for sm_107a against 35 for sm_100a and sm_103a. The toolkit floor is CUDA 13.4.cake_gdn_fla_cp_detbuilds sm_107a from its sm_103a records, and its sm_107a SASS was compared only against the previous compile flags. Validation of this tree on R200 is pending.B = 1, chunk size 64; no initial or final state under any context parallelism. A KCP call routed to the FLA-CP kernels needs a contiguous partition whoseworld_sizedivides the token count.prebuilddoes not cover FLA's Triton CP pre-processing.🤖 Generated with Claude Code