Skip to content

[KDA] Warn on kwargs silently dropped by chunk_kda and fused_recurrent_kda - #1122

Open
EastZeus wants to merge 1 commit into
fla-org:mainfrom
EastZeus:kda-warn-unconsumed-kwargs
Open

[KDA] Warn on kwargs silently dropped by chunk_kda and fused_recurrent_kda#1122
EastZeus wants to merge 1 commit into
fla-org:mainfrom
EastZeus:kda-warn-unconsumed-kwargs

Conversation

@EastZeus

@EastZeus EastZeus commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes #1119.

chunk_kda and fused_recurrent_kda keep a trailing **kwargs so 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 passes use_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_kwargs helper to fla.utils and calls it from both public KDA entry points after all known keys have been popped. Any leftover key now raises a UserWarning naming the ignored arguments and pointing at a likely version mismatch.

Design notes:

  • A warning rather than a 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.
  • Documented pass-through keys stay silent: A_log, dt_bias, chunk_size, and the deprecated transpose_state_layout (which keeps its existing DeprecationWarning).
  • The helper is generic on purpose; if this shape looks right, follow-up PRs can wire it into the other op families that share the same bare **kwargs pattern (gdn2, gated_delta_rule, gla, rwkv7, etc.).

Blast radius: no numerical or API changes. The only observable difference is a new UserWarning on calls that pass arguments the installed version does not consume, which previously did nothing.

Test plan

  • Added test_unconsumed_kwargs_warning to tests/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.py selects tests/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.
  • Local machine is CPU-only (no triton), so the kernel-path tests need CI hardware. Verified locally: helper logic exercised directly (warns on unknown keys, sorted key list in message, silent when consumed or empty), ruff check clean 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

  • I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
  • I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
  • This is not a minor/cosmetic-only PR (typo, formatting, style-only tweaks).
  • Dependent tests pass locally or in CI; new behavior is covered by tests where applicable. (New behavior is covered by the added test; dependent GPU tests need CI, no GPU available locally.)
  • Kernel changes include same-hardware before/after benchmark numbers (dense + varlen where applicable). (Not a kernel change.)

…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
@zhiyuan1i zhiyuan1i added the minor Low-value-density change (typo/docs/small validation), batch-process label Aug 14, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

minor Low-value-density change (typo/docs/small validation), batch-process

Projects

None yet

2 participants