[GDN] Add FlashQLA backend dispatch by Erix025 · Pull Request #998 · fla-org/flash-linear-attention

Erix025 · 2026-07-01T03:52:52Z

Summary

Add FlashQLA as a backend for chunk_gated_delta_rule. FlashQLA is a TileLang-based fused forward/backward kernel for the Gated Delta Rule, developed by Qwen Team (Alibaba Cloud). It replaces the multi-kernel Triton path on supported hardware (SM90/SM100), providing both forward and backward pass support including varlen, GVA, and state_v_first. It also introduce intra-card CP to achieve better performance.

Changes:

Register FlashQLABackend in the GDR backend dispatch system (priority=3)
Verifier rejects unsupported configurations: non-Hopper/SM100 GPU, K != 128, V != 128, use_gate_in_kernel, use_beta_sigmoid_in_kernel, allow_neg_eigval, inter-card CP
Export IS_NVIDIA_SM100 from fla.utils for the compute capability check
Add correctness tests (dense, varlen, state_v_first, GVA) and verifier rejection tests

Files changed:

File	Change
`fla/ops/gated_delta_rule/backends/__init__.py`	New — register `FlashQLABackend` in `gdr_registry`
`fla/ops/gated_delta_rule/backends/flashqla.py`	New — `FlashQLABackend` with verifier + dispatch
`fla/ops/gated_delta_rule/chunk.py`	Add `@dispatch('gated_delta_rule')` decorator
`fla/utils/_device.py`	Add `IS_NVIDIA_SM100`
`fla/utils/__init__.py`	Export `IS_NVIDIA_SM100`
`tests/ops/test_gdn.py`	Add FlashQLA correctness + verifier rejection tests

Test plan

test_flashqla_chunk — dense fwd+bwd, 3 configs (B/T/H vary, D=128), ratio ≤ 0.008
test_flashqla_chunk_varlen — variable-length sequences, 3 cu_seqlens configs
test_flashqla_chunk_state_v_first — state_v_first=True vs False, bit-exact (ratio=0)
test_flashqla_chunk_gva — grouped value attention (H<HV), 2 configs
test_flashqla_verifier_rejects — all 6 rejection branches + acceptance path
Ruff lint passes on all changed files
All 11 tests pass on NVIDIA H200, verified stable across multiple runs

pytest tests/ops/test_gdn.py -k "flashqla" -v

Note

FlashQLA requires the flash_qla package (QwenLM/FlashQLA). Tests skip gracefully when the package is not installed.
Disable with FLA_FLASH_QLA=0.

zhiyuan1i

Thanks for the PR! The backend dispatch looks clean. A few things to confirm (inline):

zhiyuan1i · 2026-07-02T07:10:13Z

+        if q.shape[-1] != 128:
+            return False, f"FlashQLA requires K=128, got {q.shape[-1]}"
+        if v.shape[-1] != 128:
+            return False, f"FlashQLA requires V=128, got {v.shape[-1]}"


The verifier doesn't check dtype. The default Triton path supports fp32, but flash_qla asserts q.dtype != torch.float32, so with flash_qla installed an fp32 call raises AssertionError instead of falling back to Triton. Should the verifier reject fp32 (and any other unsupported dtype)?

Thanks, fixed. The verifier now rejects unsupported input dtypes and only dispatches FlashQLA for fp16/bf16 q/k/v with matching dtype; fp32 now falls back to the Triton path.

zhiyuan1i · 2026-07-02T07:10:13Z

+    ):
+        import flash_qla
+
+        return flash_qla.chunk_gated_delta_rule(


The default Triton path returns final_state in fp32 (chunk_delta_h.py allocates it as torch.float32), but flash_qla returns it in the input dtype (bf16). This silently changes behavior once the backend is enabled. Is this intended, or should final_state be cast to fp32 to match?

I checked this again and the returned final_state is already fp32 on my side. Could you point me to the exact case or line where you observed it returning bf16? I may be missing a path, but the current output I see matches the Triton fp32 behavior.

zhiyuan1i · 2026-07-02T07:10:14Z

+            output_final_state=output_final_state,
+            use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
+            state_v_first=state_v_first,
+            cu_seqlens=cu_seqlens,


The PR description mentions intra-card CP for better performance, but auto_cp=True is never passed (FlashQLA defaults it to False), so CP isn't actually active through FLA. Should it be enabled, or should the description be updated?

Fixed. Now we pass auto_cp=True from the FLA wrapper, and we also updated FlashQLA so intra-card auto CP is enabled by default there as well. The PR description now matches the actual dispatch behavior.

zhiyuan1i · 2026-07-02T07:10:14Z

+)
+def test_flashqla_chunk(B, T, H, D, monkeypatch):
+    torch.manual_seed(42)
+    dtype = torch.bfloat16


All FlashQLA tests use bf16; fp16 isn't covered. If fp16 is supported, could you add a test? If not, the verifier should reject it.

Fixed. I added fp16 coverage to the FlashQLA dtype test and kept bf16 coverage as well, so both supported input dtypes now run through fwd/bwd correctness checks.

…ence checks to exact equality

Erix025 · 2026-07-02T13:15:17Z

Thanks for the PR! The backend dispatch looks clean. A few things to confirm (inline):

Thanks for the careful review. We addressed the verifier dtype checks, fp16 test coverage, state_v_first exactness check, and the auto_cp path. For final_state, I rechecked and it appears to already be fp32; I left a question on that thread to confirm which case showed bf16.

- Rename backend modules to flash_qla.py / flash_kda.py; align backend_type with package_name - Verifier: read use_gate_in_kernel from kwargs (arg order matches the op), reject deprecated transpose_state_layout - Make the verifier rejection test hardware-independent via monkeypatch - Fold dtype/GVA cases into a single parametrized test_flash_qla_chunk

Erix025 added 2 commits June 30, 2026 17:54

[GDN] Add FlashQLA backend dispatch

be69fc3

style: fix autopep8 formatting

a673307

zhiyuan1i reviewed Jul 2, 2026

View reviewed changes

yzhangcs and others added 4 commits July 2, 2026 15:58

Merge branch 'main' into dev/flashqla

4936230

add dtype verifier for flashqla backend

b51cb79

test(gdn): add FlashQLA dtype coverage, tighten state_v_first equival…

f8736c2

…ence checks to exact equality

enable auto cp by default in flashqla backend

29ec697

yzhangcs added 4 commits July 7, 2026 13:25

Merge branch 'main' into dev/flashqla

822ddb2

[Utils] Wrap over-length lines in _device.py

7b8d6ae

[README] Add FlashQLA backend news entry

7e6b014

yzhangcs approved these changes Jul 7, 2026

View reviewed changes

yzhangcs merged commit 77b509f into fla-org:main Jul 7, 2026
7 of 8 checks passed

xy200303 mentioned this pull request Aug 8, 2026

[Fix] Restore backend dispatch under torch.compiler.disable #1111

Merged

5 tasks done

Navigation Menu

Repository navigation

[GDN] Add FlashQLA backend dispatch - #998

[GDN] Add FlashQLA backend dispatch#998
yzhangcs merged 10 commits into
fla-org:mainfrom
Erix025:dev/flashqla

Erix025 commented Jul 1, 2026

Uh oh!

zhiyuan1i left a comment

Uh oh!

zhiyuan1i Jul 2, 2026

Uh oh!

Erix025 Jul 2, 2026

Uh oh!

zhiyuan1i Jul 2, 2026

Uh oh!

Erix025 Jul 2, 2026

Uh oh!

zhiyuan1i Jul 2, 2026

Uh oh!

Erix025 Jul 2, 2026

Uh oh!

zhiyuan1i Jul 2, 2026

Uh oh!

Erix025 Jul 2, 2026

Uh oh!

Uh oh!

Erix025 commented Jul 2, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

Erix025 commented Jul 1, 2026

Summary

Test plan

Note

Uh oh!

zhiyuan1i 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!

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!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Erix025 commented Jul 2, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants