fla: sync with flash-linear-attention main by kashif · Pull Request #1141 · huggingface/kernels-community

kashif · 2026-09-07T13:22:54Z

Important: Read before submitting

New contributors: Please open or comment on an issue before submitting a
PR to discuss the change you'd like to make. This helps us align on approach and
avoids wasted effort on changes we may not be able to merge. Please only create
a PR once one of the project maintainers agrees on your outlined approach.

PRs of contributors who are not vouched for are automatically closed. Regular
contributors are added to the vouch list.

The kernels-community repository is for:

Kernels developed by Hugging Face and partners.

Kernels developed by third parties, that have not been 'kernelized' yet, but
are used by Hugging Face projects such as diffusers and transformers.

We cannot accept PRs for the following:

New kernels that are not in categories (1) and (2).

Changes to kernels in category (2), these should be submitted upstream.

For, LLM-generated changes, we prefer that you write your prompt in an issue
over a PR with LLM-generated changes.

Related issue

Closes #

What does this PR do?

Syncs the vendored fla kernel with the current flash-linear-attention main (e52dbc0).

Motivation

fla was a couple months stale. Routine refresh.

Changes

Updated fla/torch-ext/fla/{ops,modules,utils} from upstream, new absolute imports rewritten to relative as usual.
New op families came along: atk, momentum_delta_rule, precond_gated_delta_rule, precond_kda, and a few more.
deltaformer needs pad_input/unpad_input from fla/layers/utils.py upstream (not vendored here) — originally pulled just those two functions into fla/torch-ext/fla/utils/padding.py; rebasing onto main picked up fla: fix relative import depths in the vendored source and add smoke tests #1129's fla/torch-ext/fla/utils/layers.py (a full vendored copy, same idea, landed independently while this was open), so dropped padding.py and pointed deltaformer at that instead.
__init__.py now sets __version__ (upstream's _decorators.py needs it), kept everything else about it and layers.py as-is.
Bumped build.toml version 1 → 2 (on top of main's xpu backend addition), check_public_api.py flagged new optional kwargs on chunk_gated_delta_rule/chunk_kda.
Dropped triton_ascend backends everywhere (NPU-only, build.toml declares backends = ["cuda", "xpu"], no NPU) and tilelang backends that were pure speed optimizations (rwkv6, dplr, kda, and now gdn2 too — upstream added a triton_ascend gdn2 backend in the same window). Kept ops/common/backends/tilelang (+ its attn/common registrations) since it's the only fallback for a real Hopper correctness bug on Triton 3.4.0-3.7.0, not just perf (see the RuntimeError in chunk_o.py, upstream metal-flash-sdpa: add or update card #640).
Also swept the remaining scattered NPU-only code: ascend_compile_kwargs() and a triton extra.cann→extra.ascend alias in utils/_compat.py, and an Ascend grid-size workaround in modules/token_shift.py. Left the plain IS_NPU/IS_AMD/IS_INTEL device-platform flags alone — shared, cheap, zero-dependency, not NPU-specific code.
Rebased onto main to pick up fla: fix relative import depths in the vendored source and add smoke tests #1129 (relative-import fixes + fla/tests/test_fla.py, both kept as-is) and fla: add XPU support and optimizations #1145 (XPU support — hand-merged its IS_INTEL additions in chunk_o.py/chunk_delta_h.py/gated_delta_rule/{chunk_fwd,wy_fast}.py together with this branch's own upstream sync of those same files).
Picked up the small amount of upstream flash-linear-attention movement since the last sync (516143e→e52dbc0): gdn2 gained @dispatch decorators (its new triton_ascend backend isn't vendored, same as everywhere else — dispatch() no-ops cleanly with no registry present), and solve_tril.py/utils/_device.py/utils/__init__.py picked up real AMD gfx1250 TMA support.

Testing

Imported every submodule of the built package. Everything imports cleanly except two pre-existing issues, neither introduced here (reproduced identically against flash-linear-attention itself): a missing/mismatched tilelang install, and a real upstream bug in ops/delta_rule/parallel.py (imports a name wy_fast.py doesn't define).
Re-ran check_public_api.py after every change — same two additive kwarg changes throughout, nothing new.

Checklist

This PR is linked to an issue that was discussed and approved
I have tested these changes locally
New/changed functionality has test coverage
LLM disclosure:
- I did not use an LLM to create this PR.
- I used and LLM for assistance while creating this PR.
- This PR was mostly or completely generated by an LLM.

🤖 Generated with Claude Code

vasqu

Waitng on #1129 first so we can make sure we have tests for gdn and kda

I guess this sync would also be mainly for kda and gdn? I recall some support for context parallel, can we check in layers.py as well to sync the compatibility more concretely

Update fla/torch-ext/fla/{ops,modules,utils} to fla-org/flash-linear-attention@516143e. Adds new op families (atk, momentum_delta_rule, precond_gated_delta_rule, precond_kda, and more) and picks up upstream fixes since the last sync. - Rewrote new/changed absolute fla.* imports to relative, as before. - Kept the minimal hand-maintained __init__.py/layers.py (no fla.layers/fla.models vendoring); __init__.py now sets __version__, which utils/_decorators.py needs. - Vendored pad_input/unpad_input into a new fla/torch-ext/fla/utils/padding.py so fla.ops.deltaformer works without pulling in fla/layers. - Bumped build.toml version 1 -> 2: chunk_gated_delta_rule/chunk_kda gained new optional kwargs per check_public_api.py. - Dropped triton_ascend backends everywhere (NPU-only, build.toml only declares backends = ["cuda"]) and tilelang backends that were pure speed optimizations (rwkv6, dplr, kda). Kept ops/common/backends/tilelang + its attn/common registrations: it's the only fallback for a real Hopper correctness bug on Triton 3.4.0-3.7.0 (see chunk_o.py's RuntimeError, upstream huggingface#640), not just perf. - Also dropped the remaining scattered NPU-only code paths: ascend_compile_kwargs() and the triton.language.extra.cann->ascend aliasing in utils/_compat.py, and the Ascend grid-size workaround in modules/token_shift.py. Left the plain IS_NPU/IS_AMD/ IS_INTEL device-platform flags in utils/_device.py alone - they're a shared, cheap, zero-dependency family, not NPU-specific code.

gdn2 gains dispatch() decorators on chunk_gdn2_bwd_wy_dqkg_fused/chunk_gdn2_fwd_intra (upstream added a triton_ascend gdn2 backend to hang off them; not vendored here, same as every other triton_ascend backend, and dispatch('gdn2') no-ops cleanly with no registry present). solve_tril.py and utils/_device.py+__init__.py pick up AMD gfx1250 TMA support (IS_AMD_TMA_ARCH, get_device_arch) - real, not NPU-specific. Skipped ascend_ub_manager.py changes (file doesn't exist here) and the new gdn2/backends/triton_ascend/ directory, consistent with prior NPU-code removal.

check_kernel_test_versions.py expects tests to reference build.toml's declared version; bumped it to 2 there but missed the test file itself.

e52dbc0e..a7a32d8b: fixes dx/dw/db being written/counted for rows outside this program's tile when BT > BS (upstream huggingface#1249). Skipped the other two upstream commits in this window: one touches fla/layers/comba.py (not vendored), the other touches ops/kda/backends/triton_ascend (not vendored, same as always).

danieldk · 2026-10-20T14:27:09Z

What is the status of this PR?

vasqu · 2026-10-20T14:31:07Z

More like WIP, I'd rather make my PR be merged so we have a compile baseline and then sync

vasqu

Ok we need this, best to even update more based on more recent fla main?

Lmk if you're still up for it otherwise I would take over

# Conflicts: # fla/build.toml

github-actions · 2026-10-22T11:12:10Z

This PR changes source code for the following kernel(s): fla.

If you have the kernel-bot build permissions, then trigger /kernel-bot build fla yourself. Otherwise, a maintainer will review the PR and will take an appropriate action. Thanks for your patience.

github-actions · 2026-10-22T11:17:05Z

Build request processed.

Command: /kernel-bot build weathernext2-banded-attention
Mode: build only
Target branch: pr-1141
PR head SHA: 33c7d1f5acdadbe29ce2bf6ad8c3718ce2cd86bf
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

weathernext2-banded-attention (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39769417906

kashif · 2026-10-22T11:17:53Z

/kernel-bot build fla

github-actions · 2026-10-22T11:18:07Z

Build request processed.

Command: /kernel-bot build fla
Mode: build only
Target branch: pr-1141
PR head SHA: 33c7d1f5acdadbe29ce2bf6ad8c3718ce2cd86bf
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

fla (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39769542588

vasqu · 2026-10-22T17:52:06Z

/kernel-bot build-and-stage fla

github-actions · 2026-10-22T17:52:23Z

Build request processed.

Command: /kernel-bot build-and-stage fla
Mode: build and stage
Target branch: pr-1141
PR head SHA: 5abf53015f013e432619300d793babdacd413afb
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

fla (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39826454523

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/fla

drbh

nice, lgtm!

approving. @vasqu and @kashif please feel free to merge whenever

kashif · 2026-10-22T19:06:23Z

/kernel-bot build-and-stage fla

github-actions · 2026-10-22T19:06:40Z

Build request processed.

Command: /kernel-bot build-and-stage fla
Mode: build and stage
Target branch: pr-1141
PR head SHA: 2a958266f6131ac655b729a4f1b46eef1017c776
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

fla (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39837224246

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/fla

kashif · 2026-10-22T19:19:46Z

/kernel-bot build-and-stage fla

github-actions · 2026-10-22T19:20:01Z

Build request processed.

Command: /kernel-bot build-and-stage fla
Mode: build and stage
Target branch: pr-1141
PR head SHA: 6c7f8c4975ecacbdc49f685d03d1d27fcc3b4f6c
Workflows: build.yaml, build-mac.yaml, build-windows.yaml

Dispatched (1):

fla (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39839144931

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-staging/fla

kashif · 2026-10-23T11:21:48Z

/kernel-bot merge-and-upload fla

github-actions · 2026-10-23T11:22:03Z

I can only run /kernel-bot merge-and-upload for users with write or admin repository permission.

danieldk · 2026-10-23T11:31:07Z

/kernel-bot merge-and-upload fla

github-actions · 2026-10-23T11:31:25Z

Build request processed.

Command: /kernel-bot merge-and-upload fla
Mode: merge, build and upload
Target branch: ``
PR head SHA: 6c7f8c4975ecacbdc49f685d03d1d27fcc3b4f6c
Workflows: `build.yaml, build-mac.yaml, build-windows.yaml`

Merge result: Pull Request successfully merged

Dispatched (1):

fla (build.yaml): https://github.com/huggingface/kernels-community/actions/runs/39946681306

Hub uploads:

Hub repo: https://huggingface.co/kernels/kernels-community/fla

github-actions Bot added the chore Version bumps, releases, misc maintenance label Sep 7, 2026

kashif force-pushed the sync-fla-9d981ff branch 2 times, most recently from cd74c4e to 92921d4 Compare September 7, 2026 13:42

vasqu reviewed Sep 10, 2026

View reviewed changes

kashif force-pushed the sync-fla-9d981ff branch from 92921d4 to 4b11586 Compare September 11, 2026 15:53

github-actions Bot added the needs-rebase Has merge conflicts with the base branch label Sep 11, 2026

kashif added 2 commits September 17, 2026 09:54

kashif force-pushed the sync-fla-9d981ff branch from 4b11586 to d2408df Compare September 17, 2026 07:57

github-actions Bot removed the needs-rebase Has merge conflicts with the base branch label Sep 17, 2026

kashif added 3 commits September 20, 2026 19:19

Merge branch 'main' into sync-fla-9d981ff

a345660

fla: pin test_fla.py to the bumped kernel version

0635256

check_kernel_test_versions.py expects tests to reference build.toml's declared version; bumped it to 2 there but missed the test file itself.

vasqu reviewed Oct 21, 2026

View reviewed changes

kashif added 2 commits October 22, 2026 13:02

Merge remote-tracking branch 'origin/main' into sync-fla-9d981ff

9962912

# Conflicts: # fla/build.toml

fla: sync with upstream main (12cb9d40), add upstream tests

33c7d1f

kashif and others added 5 commits October 22, 2026 13:52

fla: drop upstream tests, keep the small gdn/kda check

c5ffc53

fla: drop rwkv7 markdown

e787c64

fla: drop docs from torch-ext

f286bea

make compilable

d484b9e

more of those

39972a2

vasqu mentioned this pull request Oct 22, 2026

fla: make compile compatible after sync #1262

Closed

vasqu added 2 commits October 22, 2026 19:45

last fixes

6254d85

fixup

5abf530

drbh previously approved these changes Oct 22, 2026

View reviewed changes

fla: keep gla output scale in fp32 for rocm compile

2a95826

kashif dismissed drbh’s stale review via 2a95826 October 22, 2026 19:05

kashif added 2 commits October 22, 2026 21:19

fla: build for rocm

787e577

Merge branch 'main' into sync-fla-9d981ff

6c7f8c4

drbh approved these changes Oct 22, 2026

View reviewed changes

huggingface deleted a comment from github-actions Bot Oct 23, 2026

github-actions Bot merged commit faa8674 into huggingface:main Oct 23, 2026
9 checks passed

kashif deleted the sync-fla-9d981ff branch October 23, 2026 11:32

Conversation

kashif commented Sep 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Important: Read before submitting

Related issue

What does this PR do?

Motivation

Changes

Testing

Checklist

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

danieldk commented Oct 20, 2026

Uh oh!

vasqu commented Oct 20, 2026

Uh oh!

vasqu left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Oct 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

kashif commented Oct 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

vasqu commented Oct 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

drbh left a comment

Choose a reason for hiding this comment

Uh oh!

kashif commented Oct 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

kashif commented Oct 22, 2026

Uh oh!

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

Uh oh!

Uh oh!

kashif commented Oct 23, 2026

Uh oh!

github-actions Bot commented Oct 23, 2026

Uh oh!

danieldk commented Oct 23, 2026

Uh oh!

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

Uh oh!

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

kashif commented Sep 7, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 22, 2026 •

edited

Loading

github-actions Bot commented Oct 23, 2026 •

edited

Loading