Repository navigation
refactor(cake_kda): give single-rank and context-parallel KDA training deterministic names - #8054
Merged
Merged
Conversation
Contributor
|
No actionable comments were generated in the recent review.
 |
| Layer / File(s) | Summary |
|---|---|
Public APIs and single-rank package setup flashinfer/cake_kda_det_train.py, flashinfer/cake_kda_det_train_kcp.py, flashinfer/experimental/cake_kda_det_train/*, benchmarks/bench_cake_kda_det_train*.py |
Adds public wrappers for deterministic forward, backward, autograd, workspace, arena, partition, context, and byte-query operations. The single-rank package exports backend metadata and manifest-based library loading. Benchmarks and documentation use the renamed APIs. |
Single-rank generated kernels flashinfer/experimental/cake_kda_det_train/csrc/* |
Renames generated kernel and binding symbols to the deterministic-training names. Some kernels also update state transfers, barrier-stage indexing, or generated temporary-variable blocks. |
KCP package, build, and generated kernels flashinfer/experimental/cake_kda_det_train_kcp/* |
Adds the deterministic KCP package, build manifests, loader, lazy exports, and host-side tensor checks. Generated bindings use deterministic KCP symbols, and TMA validation reports per-axis extent errors. |
Compatibility, packaging, and validation flashinfer/experimental/cake_kda_train/*, flashinfer/kda_train.py, pyproject.toml, scripts/check_cake_build_manifests.py, tests/* |
Routes previous module paths and names to deterministic implementations. Updates package data, manifest checks, benchmarks, test helpers, and API, build, and KCP tests. |
Estimated code review effort: 5 (Critical) | ~90 minutes
Change: Refactor
Merge Risk: 
024ad
-
This change renames the Cake KDA training packages and keeps the previous names as aliases. No concrete defect remains in the reviewed changes, so it is ready to merge once the requested CI test runs pass.
Pre-merge checks: 
-
Check name Status Explanation Resolution Description check The description supplies a clear summary of the change: the deterministic renames for the single-rank and context-parallel (KCP) KDA training variants, the new package and module paths, the retained l… Complete the following before review: (1) Add the related issue or PR links in the Related Issues section, including the superseded #8082and the tracking issue for PR#8054if one exists. (2) Check the pre-commit items or state that they w…Docstring Coverage Docstring coverage is 30.43% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta… Write docstrings for the functions missing them to satisfy the coverage threshold. -
Check name Status Explanation Linked Issues check Check skipped because no linked issues were found for this pull request. Out of Scope Changes check Check skipped because no linked issues were found for this pull request. Title check The title clearly identifies the main change: renaming single-rank and context-parallel KDA training APIs to deterministic names.
Full details: Description check
-
Explanation
The description supplies a clear summary of the change: the deterministic renames for the single-rank and context-parallel (KCP) KDA training variants, the new package and module paths, the retained legacy aliases, the kernel library renames, and benchmark tables for speed. It also includes a Coverage statement naming the supported architectures (sm_100a, sm_103a, sm_107a). Parts of the template are not filled in: the Related Issues section has no link (it refers to
#8082as superseded but gives no issue number), the pre-commit and test checklist items are unchecked, there is no statement of test results or CI targets, the reason for full-scope CI is absent, and the experimental-track section is left as template text. The experimental-tests fence still contains only the template comments, even though the PR touches code under flashinfer/experimental/, which the template requires to be filled in. The PR objectives also list a tracking issue, owner, and graduation plan as required for experimental PRs, and the description does not mention them.Resolution
Complete the following before review: (1) Add the related issue or PR links in the Related Issues section, including the superseded
#8082and the tracking issue for PR#8054if one exists. (2) Check the pre-commit items or state that they were run, and state the test results. (3) Mark whether this PR requires full-scope CI; if so, give the reason; otherwise list suggested test targets. (4) Because the PR changes code under flashinfer/experimental/, complete the Experimental Track section: confirm the checklist items (tracking issue with owner, reason, graduation plan; thin core entry point; nothing registered in flashinfer/aot.py; no auto-backend reachability without the environment variable). (5) Replace the commented placeholder in the experimental-tests fence with the real test targets under tests/experimental/ and keep the fence tag unchanged. (6) Optionally add reviewer notes on the regenerated kernel sources, since the summary shows several regenerated files with changed code bodies, not only renames.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 30.43% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 69 functions across 47 files. (150 skipped: 8 unsupported, 3 too large, 139 over the file limit.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
-
Finishing Touches
-
Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
4 of 10 tasks
yyihuang
force-pushed
the
kda-kcp-det-names
branch
from
October 30, 2026 22:56
8f7cec7 to
70623fd
Compare
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_kda_det_train.py tests/experimental/test_cake_kda_det_train_api.py tests/experimental/test_cake_kda_det_train_kcp.py tests/experimental/test_cake_kda_det_train_kcp_api.py tests/experimental/test_cake_kda_det_train_kcp_build.py tests/experimental/test_cake_kda_det_train_kcp_dist.py tests/experimental/test_cake_kda_det_train_kcp_fsdp_walk.py tests/experimental/test_cake_kda_det_train_kcp_shared_ws.py tests/experimental/test_cake_kda_det_train_kcp_upload.py tests/experimental/test_cake_kda_det_train_names.py tests/experimental/test_cake_gqa_train.py |
Closed
6 tasks
flashinfer/experimental/cake_kda_train moves unchanged to flashinfer/experimental/cake_kda_det_train: the host plan, the registries, the loader, the KCP modules, build.json, build.test.json and the 165 generated sources. The packed KDA chunked training path is deterministic (the same inputs, hardware and plan give bitwise the same outputs and gradients), so the package carries the det name. The next commits add the compatibility package at the previous path, the det entry points and the sources regenerated under the det library names. Co-Authored-By: Claude <noreply@anthropic.com>
flashinfer.experimental.cake_kda_train is now a compatibility package. Its names are the det package's objects: its __all__ is the list it exported before the move, so a star import binds the same names as before, and any other attribute resolves on the det package at first access. Its submodules cake_backend, cake_build, cake_jit, cake_kcp, cake_kcp_jit and cake_kcp_plan are the det package's module objects, so the registries, the workspace caches and the loaded libraries exist once. The compatibility package holds no sources and no manifest, and importing it does not import the KCP host module. Co-Authored-By: Claude <noreply@anthropic.com>
…points The single-rank entry points move to flashinfer.cake_kda_det_train: chunk_kda_det_fwd, chunk_kda_det_bwd, chunk_kda_det, kda_det_train_arena, kda_det_train_workspace, kda_det_train_workspace_bytes, kda_det_train_saved_bytes, is_kda_det_train_supported and DEFAULT_SCALE. The backend names are ChunkKDADetFunction, KdaDetArena, KdaDetChunkSaved, KdaDetWorkspace, chunk_kda_det_fwd, chunk_kda_det_bwd and chunk_kda_det. flashinfer.kda_train keeps the context-parallel entry points unchanged and binds the previous single-rank names to the same objects. cake_backend and the package __init__ bind the previous backend names to the same objects (not in __all__), so the KCP modules and existing callers keep working and isinstance holds across the two spellings. The doc_lengths error names the det package. The package README describes the det names. The host plan of cake_jit.py names the 65 single-rank programs by their det libraries, cake_kda_det_train_<digest> (the same digests): the export reads these names and refuses a record whose library is not the det program name, so the next commit regenerates the programs under them. Co-Authored-By: Claude <noreply@anthropic.com>
…he det library names cake_jit.py, build.json and csrc/ come from a fresh export of the package, written into the moved directory. The 65 single-rank programs keep their digests and become the libraries cake_kda_det_train_<digest>: their kernel and binding sources, kernel symbols, FFI symbols, build.json entries and aot/ paths carry the det name. The 17 KCP libraries, their sources and their 33 build.json entries, build.test.json and cake_kcp_jit.py are unchanged. build.json names the package flashinfer.experimental.cake_kda_det_train. The export also writes 20 kernel sources whose local variable numbering differs from the previous export (for example j_2 -> j_0); their programs, digests and launch records are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
build.test.json carries the package of the build.json beside it, which cake_kcp_jit checks before it serves a test-only library. Its package field becomes flashinfer.experimental.cake_kda_det_train, the value the KCP export writes for the moved directory; its 18 entries are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
…hmark tests/experimental/test_cake_kda_det_train.py (the backend suite, through the cake_kda_det_train key of tests/test_helpers/cake_package) checks the cake_kda_det_train_<digest> libraries and the det package of build.json. tests/experimental/test_cake_kda_det_train_api.py covers the det entry points of flashinfer.cake_kda_det_train (its forwarding checks patch the det backend names) and the KCP entry points of flashinfer.kda_train. The reference helper is tests/test_helpers/cake_kda_det_train_reference.py, with its previous path as a compatibility import, and the benchmark is benchmarks/bench_cake_kda_det_train.py. tests/experimental/test_cake_kda_det_train_names.py checks that every previous name (facade, backend, package, compatibility package and its submodules, test helpers) is the det object, that the det __all__ lists only the det names, that the compatibility package's __all__ and star import keep the list before the move, that the manifests name the det package and the 65 single-rank libraries, and that no facade or package import loads the KCP host module. The GQA suite reads the shared loader from the det package directory, and the KCP build suite expects the det package in both manifests. Co-Authored-By: Claude <noreply@anthropic.com>
pyproject.toml ships csrc/ and build.json of flashinfer.experimental.cake_kda_det_train, the build manifest check reads that package, the pre-commit exclusion keeps the bytes of its generated csrc/ tree, and the labeler matches the det package directories next to the previous one. Co-Authored-By: Claude <noreply@anthropic.com>
…det_kcp paths Pure moves: cake_kcp.py, cake_kcp_jit.py, cake_kcp_plan.py and the test-only manifest build.test.json go to the new package flashinfer/experimental/cake_kda_det_train_kcp; the KCP test files and the KCP benchmark take the cake_kda_det_train_kcp names. The next commits rename the contents and keep the previous import paths. Co-Authored-By: Claude <noreply@anthropic.com>
…points The context-parallel (KCP) KDA training route takes deterministic names. flashinfer/cake_kda_det_train_kcp.py is the canonical entry module: partition_det_kcp, build_context_det_kcp, chunk_kda_det_kcp, chunk_kda_det_fwd_kcp, chunk_kda_det_bwd_kcp, kda_det_workspace_bytes_kcp, kda_det_saved_bytes_kcp and kda_det_comm_bytes_kcp, with the types ContextDetKcp, PartitionDetKcp, OptionsDetKcp, RecomputeHandleDetKcp, SavedDetKcp and WorkspaceDetKcp resolved on first access. The backend modules of flashinfer.experimental.cake_kda_det_train_kcp define the renamed objects (also CallDetKcp, FragmentDetKcp, ChunkKDADetFunctionKcp); their __qualname__ and error prefixes carry the new names. The previous names stay bound to the same objects: in flashinfer.kda_train, in cake_kcp / cake_kcp_plan (outside __all__), and at the previous module paths flashinfer.experimental.cake_kda_train.cake_kcp, cake_kcp_jit and cake_kcp_plan. The KCP section of the single-rank README moves to the new package's README; the single-rank manifest test no longer counts the KCP programs, which the next commit moves to the new package. The KCP tests and benchmark import the new paths; the KCP api cases move to test_cake_kda_det_train_kcp_api.py with alias-identity tests. Co-Authored-By: Claude <noreply@anthropic.com>
The eleven shipped context-parallel programs keep the code they ship with; only their names change. Their CUDA sources (file names, kernel symbols, binding namespaces) are written to cake_kda_det_train_kcp/csrc and their build.json entries (names, paths, sources, digests) to the package's build.json, under the cake_kda_det_train_<digest>[_kcp] names; the registry cake_kcp_jit.py uses the same names. All other bytes are unchanged. Co-Authored-By: Claude <noreply@anthropic.com>
… the single-rank package The six test-only context-parallel programs (built only by the tests, build.test.json) are regenerated under the det_kcp names in the det_kcp package; their kernel SASS is unchanged on sm_100a, sm_103a and sm_107a. The single-rank package no longer carries the context-parallel programs: their units, their build.json entries and its build.test.json are removed. The single-rank programs' units and entries are unchanged; the build.json generator record names the revision of this render. Co-Authored-By: Claude <noreply@anthropic.com>
yyihuang
marked this pull request as ready for review
October 31, 2026 02:58
yyihuang
requested review from
Anerudhan,
Jingfan-Sun,
YangXu1990uiuc,
aleozlx,
bkryu,
flashinfer-bot,
jhjpark,
jimmyzho,
kahyunnam and
yzh119
as code owners
October 31, 2026 02:58
yyihuang
force-pushed
the
kda-kcp-det-names
branch
from
October 31, 2026 03:00
70623fd to
024adad
Compare
yyihuang
disabled auto-merge
October 31, 2026 04:06
yyihuang
enabled auto-merge (squash)
October 31, 2026 04:06
yyihuang
disabled auto-merge
October 31, 2026 04:08
Signed-off-by: Yingyi Huang <yingyihuang2000@outlook.com>
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_gqa_train.py tests/experimental/test_cake_kda_det_train.py tests/experimental/test_cake_kda_det_train_api.py tests/experimental/test_cake_kda_det_train_kcp.py tests/experimental/test_cake_kda_det_train_kcp_api.py tests/experimental/test_cake_kda_det_train_kcp_build.py tests/experimental/test_cake_kda_det_train_kcp_dist.py tests/experimental/test_cake_kda_det_train_kcp_fsdp_walk.py tests/experimental/test_cake_kda_det_train_kcp_shared_ws.py tests/experimental/test_cake_kda_det_train_kcp_upload.py tests/experimental/test_cake_kda_det_train_names.py tests/test_cake_causal_conv1d_train_kcp_dist.py |
yyihuang
enabled auto-merge (squash)
October 31, 2026 04:34
Collaborator
Author
|
@flashinfer-bot run tests/experimental/test_cake_gqa_train.py tests/experimental/test_cake_kda_det_train.py tests/experimental/test_cake_kda_det_train_api.py tests/experimental/test_cake_kda_det_train_kcp.py tests/experimental/test_cake_kda_det_train_kcp_api.py tests/experimental/test_cake_kda_det_train_kcp_build.py tests/experimental/test_cake_kda_det_train_kcp_dist.py tests/experimental/test_cake_kda_det_train_kcp_fsdp_walk.py tests/experimental/test_cake_kda_det_train_kcp_shared_ws.py tests/experimental/test_cake_kda_det_train_kcp_upload.py tests/experimental/test_cake_kda_det_train_names.py tests/test_cake_causal_conv1d_train_kcp_dist.py |
yyihuang
added a commit
to yyihuang/flashinfer
that referenced
this pull request
Oct 31, 2026
yyihuang
added a commit
to yyihuang/flashinfer
that referenced
this pull request
Oct 31, 2026
yyihuang
added a commit
to yyihuang/flashinfer
that referenced
this pull request
Oct 31, 2026
yyihuang
disabled auto-merge
October 31, 2026 05:34
yyihuang
enabled auto-merge (squash)
October 31, 2026 05:34
yyihuang
disabled auto-merge
October 31, 2026 05:41
yyihuang
enabled auto-merge (squash)
October 31, 2026 05:46
yyihuang
pushed a commit
to yyihuang/flashinfer
that referenced
this pull request
Oct 31, 2026
Co-Authored-By: Claude <noreply@anthropic.com>
8 tasks done
yyihuang
added a commit
that referenced
this pull request
Oct 31, 2026
… backend (#8059) The chunked LM-head + loss backend is deterministic: for fixed inputs, hardware, chunk plan and initial accumulation buffers, its forward and backward give bitwise-identical results run to run. This PR gives it `det` names: - the package `flashinfer.experimental.cake_lm_head_loss` becomes `flashinfer.experimental.cake_lm_head_det_loss`; its generated sources, kernels and libraries are named `cake_lm_head_det_loss_*`; - the public entries become `chunked_lm_head_det_{loss,logprob,preload,support_logprob,support_kl}`, and the backend's functions and classes get the same `det` names. Compatibility: every previous name stays as an alias of the same object, including `load_cake_lm_head_loss_module` in `cake_jit`, `cake_backend` and `cake_backend_opt_in`. The previous package path stays as a shim package whose modules are the `det` package's modules. The `FLASHINFER_CAKE_LM_HEAD_LOSS_*` environment variables keep their names. Performance is unchanged: the generated programs differ from main only in their names, and their SASS is identical to main's on sm_100a and sm_103a. The backend has one data type (BF16) and covers the CE and policy losses, the log-probability entry and the support-set log-probability and KL entries; it ships on sm_100a and sm_103a. Merge after #8054, which carries the pre-commit exclusions for both package paths, and after the companion internal export change; a final build.json-only commit will record the published generator revision before merge. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> <!-- coderabbit-ai-generated:v1 --> ## Summary by CodeRabbit * **New Features** * Added deterministic, chunked language-model head loss and log-probability operations, along with candidate-restricted log-probabilities and KL calculations. * Added fast and memory-budgeted backend options, including valid-row controls and workspace reporting. * **Compatibility** * Preserved legacy API and package paths as aliases to the deterministic entry points. * **Documentation** * Clarified determinism guarantees and updated guidance for the APIs, naming conventions, and tuning options. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Yingyi Huang <averyh@nvidia.com> Co-authored-by: Claude <noreply@anthropic.com>
This was referenced Nov 5, 2026
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.
The KDA (Kimi Delta Attention) chunked training operator has two variants. Both are deterministic BF16: one is single-rank and the other is context-parallel (KCP). This PR gives both variants deterministic names. It supersedes #8082.
Only the names change in behavior. Some generated kernel sources are regenerated from the current kernel sources; the program digests are unchanged.
Single-rank. The public functions are now
chunk_kda_det,chunk_kda_det_fwdandchunk_kda_det_bwdinflashinfer.experimental.cake_kda_det_train. The previous name waschunk_kdainflashinfer.experimental.cake_kda_train.Context-parallel (KCP). The package is now
flashinfer.experimental.cake_kda_det_train_kcp, and the entry module isflashinfer.cake_kda_det_train_kcp. The new names are:chunk_kda_det_kcp,chunk_kda_det_fwd_kcp,chunk_kda_det_bwd_kcp,partition_det_kcp,build_context_det_kcpkda_det_workspace_bytes_kcp,kda_det_saved_bytes_kcp,kda_det_comm_bytes_kcpRecomputeHandleDetKcp,WorkspaceDetKcp,SavedDetKcp,OptionsDetKcp,CallDetKcpThe previous names still point to the same objects. Examples include
chunk_kda_kcp,kcp_partition,build_kcp_contextandKcpRecomputeHandleinflashinfer.kda_train. The modulesflashinfer.experimental.cake_kda_train.cake_kcp,cake_kcp_planandcake_kcp_jitalso still work.Kernel libraries are renamed to
cake_kda_det_train_<digest>. Programs used only for context parallelism also get a_kcpsuffix.Speed
Single-rank
The baseline is the same operator on main under its previous names. After/before ranges from 0.996 to 1.021, with a median of 1.000. A value below 1 means the new build is faster. In the pass column, fwd = forward, bwd = backward and ckpt = checkpointed forward+backward.
NVIDIA B200 (sm_100a), CUDA 13.4. Kernel time per call, median of 15 interleaved runs per side; kernels compiled with CUDA 13.4 (nvcc 13.4.92), torch 2.14.1 with CUDA 13.0 runtime.
NVIDIA GB300 (sm_103a), CUDA 13.4. Kernel time per call, median of 7 interleaved runs per side; kernels compiled with CUDA 13.4 (nvcc 13.4.92), torch 2.14.1 with CUDA 13.0 runtime.
sm_107a GPU (216 SMs), CUDA 13.4. Kernel time per call, median of 7 interleaved runs per side; kernels compiled with CUDA 13.4 (nvcc 13.4.92), torch 2.14.1 with CUDA 13.0 runtime.
Context-parallel (KCP)
NVIDIA GB200 (sm_100a), 4 GPUs on one node, torch 2.14.1 with CUDA 13.0 runtime. The benchmark is
benchmarks/bench_cake_kda_train_kcp.py, run with torchrun as one process per rank. Each iteration times the KCP causal conv1d and then the KCP KDA step, forward and backward, with CUDA events. The slowest rank's time counts. The baseline is the same path under the previous names. The workload is a packed seqlen of 65536 with 3 documents (15945, 21845 and 27746 tokens), 16 heads and 4096 channels. Each value is the median over interleaved runs per side, in the order base, branch, branch, base. Each run's figure is the benchmark's median over 3 groups of 8 timed iterations. After/before ranges from 0.9972 to 1.0059, so speed is unchanged within noise.contiguouszigzagzigzag_per_documentcontiguouszigzagzigzag_per_documentCoverage
The KDA family ships these two training variants, single-rank and context-parallel, on sm_100a, sm_103a and sm_107a.