Support nested FLA kernel imports for fla-core by DimensionSTP · Pull Request #48221 · huggingface/transformers

DimensionSTP · 2026-08-22T20:22:11Z

What does this PR do?

This fixes optional FLA kernel resolution when only fla-core is installed.

In the core-only layout, import fla does not necessarily attach fla.ops to the package object. The existing package-root traversal may therefore return None even though these nested kernels are directly importable:

fla.ops.gated_delta_rule.chunk_gated_delta_rule
fla.ops.kda.chunk_kda

The existing package-root resolution remains unchanged. If it returns None, this PR explicitly imports the mapped nested module as a second option.

This behavior is restricted to an explicit package whitelist containing only fla. Non-whitelisted packages retain the existing behavior, and failed nested imports continue to use the Torch fallback.

Fixes #48148

Verification

tests/kernels/test_kernels.py::TestKernelFallback: 4 passed
tests/utils/test_import_utils.py: 161 passed
Ruff lint and format checks: passed
Python compilation: passed
git diff --check: passed
CUDA smoke verification with fla-core==0.5.2:
- Gated delta rule and KDA nested kernels resolved successfully
- Forward and backward completed with finite outputs and gradients

Code Agent Policy

The Transformers repo is currently being overwhelmed by a large number of PRs and issue comments written by
code agents. These often are low-quality, or fix extremely minor issues that occur rarely or never in practice.
As a result, we're instituting a rule that first-time contributors should not use code agents to submit PRs or issues.
We'd also ask autonomous "OpenClaw"-like agents not to open any PRs or issues.

Issues/PRs from first-time contributors that violate this rule will probably just be closed without review, and we
might block you, especially if you open more than one or appear to be deliberately ignoring this. We especially do not
want new contributors to jump in on random issues to contribute an agent-written fix. This creates lots of noise
for reviewers and other users and will almost certainly get you blocked.

For more information, please read CONTRIBUTING.md.

(First-time contributors only): I confirm that this PR description and code is not written by an LLM or code agent

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline and the
Pull Request checks?
Was this discussed/approved via a Github issue or the forum? See Installed nested optional kernels silently fall back to the Torch implementation #48148 and the discussion below.
Did you make sure to update the documentation with your changes according to the guidelines?
Did you write any new necessary tests?

Who can review?

@vasqu @drbh

vasqu · 2026-08-24T07:38:53Z

Uh oh, could this actually be solved by changing

transformers/src/transformers/integrations/hub_kernels.py

Line 68 in c7cf04b

"recurrent_gated_delta_rule": "ops.gated_delta_rule",

transformers/src/transformers/integrations/hub_kernels.py

Line 209 in c7cf04b

"recurrent_gated_delta_rule": {

to fused_recurrent_gated_delta_rule (fla naming convention)

Ok no noticed that its a different issue alltogether

vasqu · 2026-08-24T07:44:37Z

I'm not sure I can follow 100%: Does this mean that the import path is different as per FLA versions or does this mean that the way the FLA package is installed may have changed the path - just would like this to be clarified so I have the correct overview

In any case open to fix this! It should definitely try to work for all kind of installs / envs

DimensionSTP · 2026-08-24T07:54:25Z

I'm not sure I can follow 100%: Does this mean that the import path is different as per FLA versions or does this mean that the way the FLA package is installed may have changed the path - just would like this to be clarified so I have the correct overview

In any case open to fix this! It should definitely try to work for all kind of installs / envs

Thanks for checking. This is not caused by a different installation layout in my reproduction.

I reproduced it with the same flash-linear-attention==0.5.0 installation while changing only the Transformers version. The function is importable from the expected nested module:

from fla.ops.gated_delta_rule import chunk_gated_delta_rule

The issue is that resolve_internal_import() starts from import fla and only traverses attributes already attached to the package object. It does not import missing nested modules. Therefore, resolution depends on whether fla.ops.gated_delta_rule happened to be imported earlier.

The proposed fallback explicitly imports the mapped nested module when root traversal returns None, making resolution independent of import order. It preserves the existing Torch fallback when that nested import is unavailable.

Different FLA versions may expose additional symbols at the package root, but the reproduced regression does not require an installation-path change.

vasqu · 2026-08-24T07:57:29Z

Wouldn't it make much more sense to adjust the behavior in resolve_internal_import then? Imo, this should take care of all that mess instead of making "patches" on functions that use this fn

DimensionSTP · 2026-08-24T08:28:48Z

Wouldn't it make much more sense to adjust the behavior in resolve_internal_import then? Imo, this should take care of all that mess instead of making "patches" on functions that use this fn

Agreed, thanks for the suggestion. I moved the nested-module resolution into resolve_internal_import() and removed the wrapper-specific fallback.

The common helper now preserves the existing package-root and attribute traversal behavior, then explicitly imports the nested parent module when needed. I also moved the regression coverage to test_import_utils.py.

I verified the full test_import_utils.py suite (163 passed) and confirmed that flash-linear-attention==0.5.0 now resolves to fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule.

DimensionSTP · 2026-08-24T08:44:10Z

Hi @vasqu, the latest CI run has a single failure in an unrelated OwlViT numerical test:

tests/models/owlvit/test_modeling_owlvit.py::OwlViTModelTest::test_eager_matches_sdpa_inference_18_bf16_pad_left_no_attn_mask_sdpa_kernels

The failed shard completed with 7,317 passed tests and this single failure. This PR only changes import_utils.py and its focused tests, and the relevant local test suite passes.

Could you please rerun the failed tests_torch [shard 4/8] job?

vasqu

Heya, I noticed that this is a versioning issue no? I can repro on fla == 0.5.0 but on the most recent fla version 0.5.2, this works without this workaround.

Tbh, this seems more like an package structure issue on FLA that was only recently solved. I'd really like to avoid making this change here as the package itself should be a properly nestable module and not only work via import module on the full path.

Tl;dr: Bump fla please, it's more of a package structure issue on fla side (solved on latest versions)

DimensionSTP · 2026-08-25T07:39:36Z

Heya, I noticed that this is a versioning issue no? I can repro on fla == 0.5.0 but on the most recent fla version 0.5.2, this works without this workaround.

Tbh, this seems more like an package structure issue on FLA that was only recently solved. I'd really like to avoid making this change here as the package itself should be a properly nestable module and not only work via import module on the full path.

Tl;dr: Bump fla please, it's more of a package structure issue on fla side (solved on latest versions)

Thanks for pointing this out. I retested this in a clean environment with both flash-linear-attention==0.5.2 and fla-core==0.5.2.

On the current unmodified Transformers main, the existing resolver correctly selects:

fla.ops.gated_delta_rule.chunk.chunk_gated_delta_rule

I also ran the Qwen3.5 gated-delta wrapper on CUDA and verified that forward, backward, and all input gradients complete with finite values.

Under the same unmodified Transformers code, FLA 0.5.0 still selects the PyTorch fallback. This confirms that the behavior was caused by the older FLA package structure and is fixed in FLA 0.5.2.

You're right that the appropriate fix is to bump FLA downstream rather than change the generic Transformers resolver. Thank you for the clarification. I'll close this PR.

DimensionSTP · 2026-09-01T16:27:39Z

Thanks @imvladikon for pointing this out. I reproduced the core-only case and reopened both this PR and #48148.

On current Transformers main with only fla-core==0.5.2 installed, import fla does not expose fla.ops. The existing resolver returns None for both:

ops.gated_delta_rule.chunk_gated_delta_rule
ops.kda.chunk_kda

As a result, use_kernel_func_from_hub_with_fallback() selects the Torch fallback even though both nested FLA functions import directly and successfully complete CUDA forward/backward with finite gradients.

The earlier 0.5.2 verification had both fla-core and flash-linear-attention installed. In that configuration, importing fla.layers transitively attaches fla.ops, masking the issue.

I agree that changing resolve_internal_import() globally may be broader than necessary. I tested a narrower decorator-local approach:

preserve the existing root resolution;
only when it returns None and internal_path is configured, explicitly import f"{package}.{internal_path}";
resolve func_name from that module;
retain the existing Torch fallback when the nested import is unavailable.

The CPU prototype selects both FLA kernels correctly, preserves package-root resolution, and retains the missing-kernel fallback.

@vasqu, would this narrower change inside use_kernel_func_from_hub_with_fallback() be an acceptable direction? I have not pushed another implementation yet, pending agreement on the preferred scope.

vasqu · 2026-09-01T16:37:04Z

It's easier for me to decide when I can see an implementation :D did the earlier comments get deleted?

Anyways, this is super annoying and makes FLA package structure really unconventional. Maybe we can just adjust the original fallback fn and make a total import as second option (only for whitelisted packages like FLA)

imvladikon · 2026-09-01T16:43:00Z

@vasqu I apologize. Yes, I deleted it because I start thinking that it's indeed a problem with the fla packaging, and not sure about a solution

vasqu · 2026-09-01T16:46:55Z

No worries, I'm also a bit baffled at their package structure tbh 😢 I mean we could allow full imports but it's also a matter of, is it really intended to only have core installed for example

DimensionSTP · 2026-09-01T16:56:31Z

No worries, I'm also a bit baffled at their package structure tbh 😢 I mean we could allow full imports but it's also a matter of, is it really intended to only have core installed for example

Yes, fla-core is intended to be usable as a standalone package for the kernel-level APIs used here.

The official FLA documentation describes the distribution split as follows:

fla-core contains the customized kernels under fla/ops, along with fla/modules and fla/utils, and has minimal dependencies.
flash-linear-attention is the extension package that adds fla/layers and fla/models and depends on fla-core.

References:

Therefore, the core-only environment is an intentional package configuration rather than a partially installed full package. It is also sufficient for the two kernel paths reproduced here: chunk_gated_delta_rule and chunk_kda.

I will update the implementation so that:

the existing package-root resolution remains unchanged;
the explicit nested-module import is used only as a second option;
this second option is restricted to an explicit package whitelist containing FLA;
non-whitelisted packages retain the current behavior;
regression tests cover the core-only FLA layout and the existing fallback behavior.

I will push the implementation and verification results to this PR for review.

vasqu · 2026-09-01T17:08:06Z

Sounds good! This FLA structure really is unconventional tho... Anyways I think it's the right approach to treat pkgs like FLA a bit more special then

DimensionSTP · 2026-09-01T17:42:59Z

Implemented and pushed in 71783437f9.

The updated implementation follows the agreed scope:

preserves the existing package-root resolution;
attempts an explicit nested-module import only as a second option;
restricts that behavior to the FLA whitelist;
leaves non-whitelisted packages unchanged;
preserves the Torch fallback when the nested module is unavailable.

The PR now changes only:

src/transformers/integrations/hub_kernels.py
tests/kernels/test_kernels.py

Verification completed:

tests/kernels/test_kernels.py::TestKernelFallback: 4 passed
tests/utils/test_import_utils.py: 161 passed
Ruff lint and format checks: passed
Python compilation: passed
git diff --check: passed
CUDA smoke tests with fla-core==0.5.2: both gated delta rule and KDA completed forward and backward with finite outputs and gradients

One CI job, examples_torch [shard 0/1], failed due to a runner/container infrastructure error rather than a test failure:

ttrpc: closed
container command terminated with exit code 137
custom container implementation failure

No test assertion or traceback related to this PR was reported. Could you please rerun that failed job?

Ready for review.

vasqu

Got some smaller comments but I think it mightve been overkill with the whitelist (sorry)

vasqu · 2026-09-02T13:14:46Z

+            if implementation is None and package in _KERNELS_EXPLICIT_IMPORT_PACKAGES and internal_path is not None:
+                internal_module = importlib.import_module(f"{package}.{internal_path}")
+                implementation = getattr(internal_module, func_name, None)


Would it make sense to change the declarations a bit so

full_path -> full_func_path

full_module_path

Suggested change

if implementation is None and package in _KERNELS_EXPLICIT_IMPORT_PACKAGES and internal_path is not None:

internal_module = importlib.import_module(f"{package}.{internal_path}")

implementation = getattr(internal_module, func_name, None)

if implementation is None and internal_path is not None:

module = importlib.import_module(full_module_path)

implementation = getattr(module, func_name, None)

vasqu · 2026-09-02T13:15:08Z

    "selective_scan_fn": "ops.selective_scan_interface",
 }

+_KERNELS_EXPLICIT_IMPORT_PACKAGES = {"fla"}


At second thought it might be a bit overkill

vasqu · 2026-09-02T13:15:37Z

@@ -56,6 +55,113 @@
    import transformers.integrations.hub_kernels as hub_kernels_pkg


+class TestKernelFallback(TestCasePlus):


Let's add the tests under the utils class no?

DimensionSTP · 2026-09-03T05:04:55Z

Addressed the latest review comments in 33f525b96e:

renamed full_path to full_func_path and introduced full_module_path;
removed the package whitelist;
kept nested-module import as a generic second-stage fallback;
moved the regression tests into the existing TestKernelUtilities class.

Verification completed:

focused regression tests: 4 passed;
Ruff lint and format checks: passed;
Python compilation and git diff --check: passed;
fla-core==0.5.2 CUDA smoke: gated-delta and KDA forward/backward completed with finite outputs, states, and gradients.

The updated commit is ready for review.

vasqu

Thanks 🫡 really nitpicky but wanna make sure this goes well as its a special case

vasqu · 2026-09-03T15:26:25Z

            module = importlib.import_module(package)
-            implementation = resolve_internal_import(module, full_path)
+            implementation = resolve_internal_import(module, full_func_path)
+            if implementation is None and internal_path is not None:


Suggested change

if implementation is None and internal_path is not None:

if implementation is None and full_module_path != package:

nit

vasqu · 2026-09-03T15:27:03Z

            module = importlib.import_module(package)
-            implementation = resolve_internal_import(module, full_path)
+            implementation = resolve_internal_import(module, full_func_path)
+            if implementation is None and internal_path is not None:


Let's add a small comment here and mention fla as example

vasqu · 2026-09-03T15:29:42Z

@@ -419,6 +418,110 @@ def test_enable_hub_kernels(self):

 @require_kernels
 class TestKernelUtilities(TestCasePlus):
+    def test_fallback_resolves_function_from_package_root(self):


2 smaller nits:

Let's add new tests at the end pls

Could we parametrize these tests a bit? Unsure if we really need separate tests for all of those

DimensionSTP · 2026-09-03T17:01:22Z

Thanks @vasqu! I addressed all of the review comments in 4c13bec07c:

simplified the nested-import guard to compare full_module_path with the package root;
added a short comment explaining the FLA case;
moved the fallback regression tests to the end of TestKernelUtilities;
parameterized the explicit and mapped nested-module cases while keeping the package-root and missing-module branches separate.

I also rebuilt the branch on top of the latest main so it includes the recent kernel utility changes.

Validation:

focused fallback regression tests: 4 passed;
Ruff lint and format checks: passed;
Python compilation and git diff --check: passed.

Ready for another review.

vasqu

Just some very last comments but we should get it merged after that, thanks a lot 🤗

vasqu · 2026-09-03T17:20:12Z

+    )
+    def test_fallback_imports_nested_module(
+        self,
+        _case_name,


Suggested change

_case_name,

case_name,

lets avoid it please

vasqu · 2026-09-03T17:20:44Z

+                internal_path="ops.kernel",
+            )(torch_function)
+
+        self.assertEqual(wrapped(), "torch")


Let's add small comments / docstrings as well just for the general what are we trying here but nice tests otherwise

DimensionSTP · 2026-09-03T17:29:23Z

Addressed the final review comments in 8f40ff6756:

renamed _case_name to case_name;
added short docstrings explaining the purpose of each fallback regression test.

Validation:

focused fallback regression tests: 4 passed;
Ruff lint and format checks: passed;
Python compilation and git diff --check: passed.

Thanks again for the review! @vasqu

vasqu · 2026-09-03T17:33:18Z

Merging now, thanks for sticking through @DimensionSTP and also @imvladikon for reopening this issue otherwise we wouldve ignored it for who knows how long

HuggingFaceDocBuilderDev · 2026-09-03T17:44:49Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-09-03T17:44:51Z

CI recap

Dashboard: View test results in Grafana
Latest run: 33784684958:2
Result: failure | Jobs: 16 | Tests: 162,287 | Failures: 0 | Duration: 14h 17m

* fix: resolve nested optional kernel imports * test: clarify nested kernel fallback coverage

DimensionSTP mentioned this pull request Aug 22, 2026

Installed nested optional kernels silently fall back to the Torch implementation #48148

Closed

2 of 4 tasks

vasqu reviewed Aug 24, 2026

View reviewed changes

vasqu reviewed Aug 25, 2026

View reviewed changes

DimensionSTP closed this Aug 25, 2026

DimensionSTP reopened this Sep 1, 2026

DimensionSTP changed the title ~~Fix nested optional kernel fallback resolution~~ Support nested FLA kernel imports for fla-core Sep 1, 2026

vasqu reviewed Sep 2, 2026

View reviewed changes

vasqu mentioned this pull request Sep 2, 2026

Kimi linear #48250

Merged

vasqu reviewed Sep 3, 2026

View reviewed changes

fix: resolve nested optional kernel imports

4c13bec

DimensionSTP force-pushed the fix/optional-kernel-nested-fallback-import branch from cdf94e6 to 4c13bec Compare September 3, 2026 16:57

vasqu approved these changes Sep 3, 2026

View reviewed changes

test: clarify nested kernel fallback coverage

8f40ff6

vasqu enabled auto-merge September 3, 2026 17:32

vasqu added this pull request to the merge queue Sep 3, 2026

vasqu removed this pull request from the merge queue due to a manual request Sep 3, 2026

vasqu merged commit c119ec3 into huggingface:main Sep 3, 2026
216 of 218 checks passed

imvladikon mentioned this pull request Sep 5, 2026

[GLM 5.3 Flash] Fix NaN gradients in chunked KDA #48455

Merged

sbucaille pushed a commit to sbucaille/transformers that referenced this pull request Sep 16, 2026

Support nested FLA kernel imports for fla-core (huggingface#48221)

3fed184

* fix: resolve nested optional kernel imports * test: clarify nested kernel fallback coverage

		@@ -56,6 +55,113 @@
		import transformers.integrations.hub_kernels as hub_kernels_pkg


		class TestKernelFallback(TestCasePlus):

	if implementation is None and internal_path is not None:
	if implementation is None and full_module_path != package:

Conversation

DimensionSTP commented Aug 22, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Verification

Code Agent Policy

Before submitting

Who can review?

Uh oh!

vasqu Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vasqu commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

DimensionSTP commented Aug 24, 2026

Uh oh!

vasqu commented Aug 24, 2026

Uh oh!

DimensionSTP commented Aug 24, 2026

Uh oh!

DimensionSTP commented Aug 24, 2026

Uh oh!

vasqu left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DimensionSTP commented Aug 25, 2026

Uh oh!

DimensionSTP commented Sep 1, 2026

Uh oh!

vasqu commented Sep 1, 2026

Uh oh!

imvladikon commented Sep 1, 2026

Uh oh!

vasqu commented Sep 1, 2026

Uh oh!

DimensionSTP commented Sep 1, 2026

Uh oh!

vasqu commented Sep 1, 2026

Uh oh!

DimensionSTP commented Sep 1, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DimensionSTP commented Sep 3, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DimensionSTP commented Sep 3, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

DimensionSTP commented Sep 3, 2026

Uh oh!

vasqu commented Sep 3, 2026

Uh oh!

HuggingFaceDocBuilderDev commented Sep 3, 2026

DimensionSTP commented Aug 22, 2026 •

edited by github-actions Bot

Loading

vasqu Aug 24, 2026 •

edited

Loading

vasqu commented Aug 24, 2026 •

edited

Loading

vasqu left a comment •

edited

Loading