Repository navigation
Conversation
Copy the FlashInfer Cake packages
flashinfer/experimental/cake_gqa_train
of flashinfer-ai/flashinfer c5db8d3ccd3fa3842812280f7f208718dc0b8610
byte-identically into miles_ext_kernels/gqa/_vendor/ with
tools/vendor_cake.py. _vendor/VENDORED.json records the provenance: the
git blob of every file and the Cake revision of each build.json.
VENDORED_FAMILIES in tests/cake_native/test_vendored.py now names gqa,
so the provenance test checks this committed _vendor/ tree and fails
when a tree is added or dropped without updating the list.
Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.gqa runs the forward and backward of FlashInfer's cake_gqa_train package (vendored separately under gqa/_vendor/cake_gqa_train, built through cake_native) for the head-dim-256 full-attention layers miles selects with --full-attention-backend cake. ops.varlen_gqa_attention is the autograd entry point over packed segments (int32 device cu_seqlens, bottom-right causal), with an SM100a / SM103a guard; ops.prebuild builds the four libraries the forward and the BF16 dK / dV backward load (fwd, bwd_pre, bwd_bf16, bwd_post) once per architecture, for the layer to call at construction. The package has one variant and no arguments. Its backward sums dQ with FP32 atomic adds in a run-dependent order, so validate_args rejects the cake backend under deterministic mode; the output, the LSE and dK / dV are bit-identical run to run. Tests: argument rules and torch-free import (CPU); the vendored launch plan, loaded by path without torch (CPU); and on SM100a / SM103a the FP64 reference at FlashInfer's budget on its suite layouts and on the TP-local head pairs 8/2, 6/1, 16/2 and 8/1 with a pad segment, run-to-run bitwise output / LSE / dK / dV, document isolation, non-reentrant and reentrant checkpointing, no synchronization in steady state, rejected operands and devices, and the prebuild hook covering every library loaded, without importing flashinfer. The package is registered in _KERNEL_PACKAGES and listed in the README kernel table. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
miles_ext_kernels.gqa now accepts compute capability 10.7: SUPPORTED_CAPABILITIES is the shared cake_native.ARCH_BY_CAPABILITY, and the stages miles runs (fwd, bwd_pre, bwd_bf16, bwd_post) launch the sm_107a libraries of the vendored cake_gqa_train package. Toolchain: device_arch, which miles calls when it builds the layer and the first call runs once per device, checks the toolkit with cake_native.check_toolkit, so on an SM107a device a toolkit older than CUDA 13.4 (nvcc 13.0 to 13.3 reject compute_107a) is refused with a RuntimeError naming the CUDA version it needs. Tests: test_gqa_gpu.py runs on SM107a too; its device test no longer expects capability 10.7 to be refused and checks the toolkit floor instead. test_gqa_cpu_plan.py checks one library per stage for sm_107a as well, against the toolkit of each architecture (CUDA 13.0 for sm_100a / sm_103a, 13.4 for sm_107a). The README gqa row lists SM107a. Co-Authored-By: Claude <noreply@anthropic.com>
Co-Authored-By: Claude <noreply@anthropic.com>
Owner
Author
|
Closed: no separate stacked review PRs; review continues on radixark#8. |
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.
Review-only stacked PR inside the fork; lands upstream via the matching radixark/miles-ext-kernels PR: radixark#8.
🤖 Generated with Claude Code