Repository navigation
Conversation
Copy the FlashInfer Cake packages
flashinfer/experimental/dense_projection_det_gemm
of flashinfer-ai/flashinfer c5db8d3ccd3fa3842812280f7f208718dc0b8610
byte-identically into miles_ext_kernels/dense_gemm/_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
dense_gemm, 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>
Adds the dense_gemm kernel package over the vendored dense projection
package (miles_ext_kernels/dense_gemm/_vendor/dense_projection_det_gemm,
brought in separately), deterministic only:
- batched_projection(x, weight, *, weight_layout): the per-head
projection x [T, H, D_in] -> [T, H, D_out] of the MLA q-absorb
("in_out", einsum thd,hdm->thm) and value ("out_in", thm,hdm->thd)
products and of DeepSeek-V4's grouped output projection, in BF16 with
FP32 accumulation, bit-identical call to call. Views the kernel reads
in place stay views; misaligned ones are copied.
- router_gating(x, weight): FP32 router logits from BF16 operands, the
input and weight gradients through the FP32 emulation GEMM, as
Megatron's router gating with --moe-router-dtype fp32 and no bias. The
forward hands K1 the weight as weight.T stored [H, E] (an N-major B
operand, a copy of E * H BF16 elements): the vendored registry serves
FP32 output with a K-major A and an N-major B ("kn"), not with a
K-major B ("kk").
- prebuild_batched_projection / prebuild_router compile the programs of
given token counts ahead through cake_native.
The package registers no arguments: --batched-projection-gemm-backend
cake and --moe-router-gemm-backend cake select it in miles, and
validate_args requires --train-backend megatron for them. Both run on
SM100a / SM103a only and are not CUDA-graph safe.
Tests: torch-free argument tests; GPU tests (SM100a / SM103a) against
FP64 for the GLM-5 / GLM-5.3 / DeepSeek-V4 projection rows and the
router's expert counts and hidden sizes, bitwise reruns, copied views,
refused shapes, prebuild coverage, and the vendored package's own MLA
and router cases. The router routing check compares every row whose k-th
/ (k+1)-th FP64 logit gap clears the tie margin with the FP64 top-8, and
needs at least 95 % of the 1001 rows clear.
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.dense_gemm now accepts compute capability 10.7: SUPPORTED_CAPABILITIES is the shared cake_native.ARCH_BY_CAPABILITY, and batched_projection, router_gating and their prebuild hooks plan and launch the sm_107a programs of the vendored dense_projection_det_gemm package. Toolchain: the per-device architecture lookup, run on the first call and by the prebuild hooks, 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_batched_projection_gpu.py and test_router_gpu.py run on SM107a too; test_router_gpu.py checks the refusal of an unsupported device and of a toolkit below the device's floor. The README dense_gemm 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#9. |
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#9.
🤖 Generated with Claude Code