refactor(kda): stop 0.7 advertising unreleased KDA surface (mark prefill wrapper experimental, trim kda_kernels.__all__) - #5040
Conversation
The wrapper has not appeared in a release and its plan-and-run shape is expected to change as the recurrent_kda surface is unified (flashinfer-ai#4936), so 0.7 should not advertise it as stable. plan and run carry @flashinfer_experimental_api, which warns once per process, and the class docstring and docs page say so.
The three *_kda_prefill_sm120 names are unreleased -- kda_kernels.__all__ grew from 6 entries to 9 since v0.6.18 and these are the delta. They are dispatch plumbing for flashinfer.kda_prefill, not surface a 0.7 cut should freeze. Every consumer reaches them as attributes or through kda_kernels.sm120_prefill, so no call site changes.
…lane The experimental policy asks for an experimental API's tests under tests/experimental/. The two tests covering the wrapper's own contract -- the plan metadata it builds and the buffers it forwards -- move there; the shared cpu_route_tensors factory moves to tests/test_helpers/ because the two lanes cannot import each other. Tests whose subject is the CuTe DSL prefill kernels stay in the stable lane even though they drive the wrapper, since it is the only planned-path entry point: moving them would take kernel numerics out of 'pytest tests/'.
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Advanced Run ID: 📒 Files selected for processing (6)
Included review availability: Your plan provides up to 8 included reviews per hour; 6 remain after this review. 📝 WalkthroughWalkthroughThe pull request marks ChangesKDA prefill API
Priority: ⬇️ Low Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to This change labels the unreleased KDA prefill wrapper as experimental, narrows package exports, and reorganizes its tests without changing kernel behavior. No current merge-blocking risk is identified. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
Full details: Docstring CoverageExplanation Docstring coverage is 9.52% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 21 functions across 5 files. (1 skipped: 1 unsupported.)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
🧪 Generate unit tests (beta)
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. Comment |
|
/bot runt tests/kda |
|
Unknown Command Command Use |
2 participants
📌 Description
The
0.7holding action from #4936 (steps 3b-i and 3c): stop0.7fromadvertising KDA surface that has never shipped and that the unification plan
expects to change. Three independent commits, no kernel logic, no behaviour
change on any released path.
Commit 1 — mark
RecurrentKDAPrefillWrapperexperimental.planandrunget
@flashinfer_experimental_api(warns once per process), the classdocstring gets a
.. warning::banner naming #4936, anddocs/api/kda_prefill.rstlabels the wrapper experimental in prose. The class,its constructor and its behaviour are untouched.
Commit 2 — keep the SM120a prefill facade out of
kda_kernels.__all__.kda_kernels.__all__grew from 6 entries to 9 sincev0.6.18, and all threeadditions are the SM120a helpers (
can_implement_kda_prefill_sm120,clear_kda_prefill_sm120_caches,run_kda_prefill_sm120). They are dispatchplumbing for
flashinfer.kda_prefill, not surface a release should freeze.__all__is now byte-identical tov0.6.18's.Commit 3 — move the wrapper's own tests to the experimental lane, per the
experimental policy. Details under "Test scope" below.
Why now
None of this is in
v0.6.18— that tag has nokda_kernels/sm120_prefill/atall, and its
kda_kernels/__init__.pyhas no SM120a entries. Both changes arefree today and cost a deprecation cycle once
0.7cuts.Why commit 2 changes no call site
Removing a name from
__all__leaves the attribute in place, and everyconsumer reaches these three by attribute or by the deeper module path:
kda_prefill.py:7042,:7135kda_kernels.can_implement_…/.run_…attributedocs/api/kda_prefill.rst:343,:369flashinfer.kda_kernels.sm120_prefill.clear_…()tests/kda/test_recurrent_kda_prefill_sm120.pyfrom …kda_kernels.sm120_prefill import …benchmarks/routines/kda.py:336sm120_prefill.run_kda_prefill_sm120(…)Nothing in the tree does
from flashinfer.kda_kernels import *, nothingasserts
__all__'s contents, and no Sphinxautomoduletargetsflashinfer.kda_kernels.clear_kda_prefill_sm120_cachesstays reachable atthe exact path the docs tell users to call, which is the deeper one.
🔍 Related Issues
recurrent_kdaAPI, and settle the unreleased KDA surface before FI v0.7 #4936 — KDA API unification audit; this is steps 3b-i and 3c.is why commit 1 marks only
RecurrentKDAPrefillWrapperand not the other fiveunreleased top-level symbols. If refactor(kda): remove the unreleased recurrent KDA training/backward family #4965 does not land before the cut, those
five need marking too and this PR should grow another commit.
🧪 Tests
pytest tests/experimental/test_kda_prefill_wrapper.py→ 2 passed.pytest tests/kda/test_recurrent_kda_prefill.py→ 247 passed, 1 skipped(B200, CC 10.0) after the relocation.
pytest tests/kda/ -k 'prefill or export or api or trace'→ 523 passed,90 skipped.
pytest tests/kda/test_recurrent_kda_prefill_sm120.py→ 50 passed, 88skipped (skips are the CC 12.0 cells; this host is CC 10.0).
pytest tests/does not collect the relocated file, and thatscripts/pr_checks/experimental_test_scope.pyaccepts the fence below.planand
runcarryis_experimental; the warning fires exactly once across twoplancalls; all three SM120a names are absent from__all__yet stillpresent as attributes;
_sm120_kda_prefill_is_eligiblestill resolves.pre-commit runclean on all six touched files.📋 Not included (deliberately)
Step 3c in #4936 also proposes giving SM120a an explicit
backend=literal andgating its automatic branch with
experimental_auto_backends_allowed(). Bothare omitted here: they change
"auto"prefill routing on CC 12.0, and per#4936 the gating mostly converts working calls into raises — the same failure
shape #5037 is fixing for decode. That belongs in its own PR, benchmarked and
validated on a CC 12.0 device, which this change does not require.
🚀 Pull Request Checklist
pre-commit.unittest, etc.).🔬 Experimental Track
flashinfer/experimental/and/or an@flashinfer_experimental_api. Tracking issue: [KDA] Unify the publicrecurrent_kdaAPI, and settle the unreleased KDA surface before FI v0.7 #4936tests/experimental/and were validated on the intended hardware; a runnable example is included.flashinfer/aot.py, and no experimental backend is reachable frombackend="auto"withoutFLASHINFER_ALLOW_EXPERIMENTAL_AUTO_BACKENDS=1. (Calling an@flashinfer_experimental_apior naming a backend explicitly is itself the opt-in and needs no environment variable.)Notes on the checklist, since this is a retrofit rather than a new feature
The policy is written for a new, self-contained feature under
flashinfer/experimental/. This is the opposite case — a marker applied toalready-merged, unreleased surface that shares kernels and tests with a stable
API — so three items deserve explanation rather than a bare tick:
Test scope. The two tests covering the wrapper's own contract (the plan
metadata it builds, the buffers it forwards) moved to
tests/experimental/test_kda_prefill_wrapper.py. Three tests that drive thewrapper but whose subject is the CuTe DSL prefill kernels stayed in
tests/kda/test_recurrent_kda_prefill.py:test_cute_dsl_checkpoints_match_cake,test_cute_dsl_padded_indexed_state_matches_cakeandtest_cute_dsl_planned_zero_length_cuda_graph_capture_and_replay. They use thewrapper because it is the only planned-path entry point, and relocating them
would take CuTe-vs-Cake numerics and graph-capture correctness out of
pytest tests/(pytest.inisetsnorecursedirs = … tests/experimental) whiledragging along that file's autouse legacy-stub fixture and ~200 lines of
reference implementation. Happy to move them too if you would rather the
experimental lane own every call site.
Runnable example. No new
examples/entry; the wrapper's usage is alreadydocumented on
docs/api/kda_prefill.rst, now labelled experimental.aot.py. This PR adds nothing toaot.py. Strictly, though, the wrapperis a thin host-side planner that forces
backend="cute-dsl"and hands off torecurrent_kda, so the kernels beneath it are the stable prefill path'sAOT-registered modules. They cannot be de-registered without de-registering
stable KDA prefill, so "experimental features are JIT-only" holds for the
wrapper's own code but not for the kernels it borrows. Flagging rather than
hiding it.
One visible side effect: the three stable-lane tests above now emit
ExperimentalWarningwhen they drive the wrapper. That is the marker workingas intended, and the suite passes with it.
Summary by CodeRabbit
Documentation
RecurrentKDAPrefillWrapperas experimental, including its packed engine CUDA Graph capture usage.API
Tests