[KDA] Warn on kwargs silently dropped by chunk_kda and fused_recurrent_kda - #1122
Open
EastZeus wants to merge 1 commit into
Open
[KDA] Warn on kwargs silently dropped by chunk_kda and fused_recurrent_kda#1122EastZeus wants to merge 1 commit into
EastZeus wants to merge 1 commit into
Conversation
…t_kda Both public KDA entry points keep a trailing **kwargs so that backend specific arguments can flow through the dispatch layer. The downside is that a misspelled argument, or a flag only known to a newer release, lands in **kwargs and is silently discarded instead of raising. This recently bit Kimi K3 users on fla 0.5.0, where the modeling code passed use_beta_sigmoid_in_kernel=True to a version that did not know the flag, producing wrong recurrences with no error or warning. Add a shared warn_unconsumed_kwargs helper to fla.utils and call it from chunk_kda and fused_recurrent_kda after all known keys have been popped, so any leftover key raises a loud UserWarning naming the ignored arguments and pointing at a likely version mismatch. Documented pass through keys (A_log, dt_bias, chunk_size, transpose_state_layout) stay silent, and the dispatch fallback path keeps working since backends can still consume their own keys. Closes fla-org#1119
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.
Summary
Fixes #1119.
chunk_kdaandfused_recurrent_kdakeep a trailing**kwargsso backend-specific arguments can flow through the dispatch layer, but that means a misspelled argument, or a flag only known to a newer release, is silently discarded instead of raising. As reported in #1119, Kimi K3's modeling code passesuse_beta_sigmoid_in_kernel=True, and on fla 0.5.0 (which predates the flag) the sigmoid was never applied: wrong recurrences, no error, no warning.This PR adds a shared
warn_unconsumed_kwargshelper tofla.utilsand calls it from both public KDA entry points after all known keys have been popped. Any leftover key now raises aUserWarningnaming the ignored arguments and pointing at a likely version mismatch.Design notes:
TypeError, because the dispatch layer legitimately lets unknown keys fall through to the default implementation when a specialized backend rejects a call (e.g. the FlashKDA path), and a hard error would break that fallback.A_log,dt_bias,chunk_size, and the deprecatedtranspose_state_layout(which keeps its existingDeprecationWarning).**kwargspattern (gdn2, gated_delta_rule, gla, rwkv7, etc.).Blast radius: no numerical or API changes. The only observable difference is a new
UserWarningon calls that pass arguments the installed version does not consume, which previously did nothing.Test plan
test_unconsumed_kwargs_warningtotests/ops/test_kda.py: asserts the warning fires for an unknown kwarg on both entry points, and that clean calls (including the documented pass-through keys) stay silent.python scripts/find_dependent_tests.py fla/ops/kda/chunk.py fla/ops/kda/fused_recurrent.py fla/utils/_decorators.pyselectstests/ops/test_kda.py,tests/ops/test_intracard_cache.py,tests/models/test_modeling_kda.py,tests/context_parallel/test_cp_kda.py, and the layer/hybrid suites.ruff checkclean with the repo config, and all touched files byte-compile.Benchmark / NCU (kernel changes only)
Neutral, no kernel changes. The helper is a set difference over an almost always empty dict, executed once per entry-point call on the host.
Breaking changes
None.
Checklist