[Refactor] Move fla to third party (#48500)

Signed-off-by: yewentao256 <[email protected]>

.buildkite/test-amd.yaml

@@ -1444,9 +1444,9 @@ steps:
   optional: true
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
-  - vllm/model_executor/layers/fla/ops/kda.py
-  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
-  - vllm/model_executor/layers/fla/ops/l2norm.py
+  - vllm/third_party/flash_linear_attention/ops/kda.py
+  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
+  - vllm/third_party/flash_linear_attention/ops/l2norm.py
   - tests/kernels/test_kda.py
   - vllm/platforms/rocm.py
   commands:
@@ -3234,7 +3234,7 @@ steps:
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py
   - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/model_executor/layers/fla/ops/
+  - vllm/third_party/flash_linear_attention/ops/
   - vllm/_aiter_ops.py
   - vllm/platforms/rocm.py
   commands:
@@ -3473,7 +3473,7 @@ steps:
   - vllm/model_executor/models/qwen3.py
   - vllm/model_executor/models/qwen3_next.py
   - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/model_executor/layers/fla/ops/
+  - vllm/third_party/flash_linear_attention/ops/
   - vllm/_aiter_ops.py
   - vllm/v1/attention/backends/triton_attn.py
   - vllm/v1/attention/backends/rocm_attn.py

.buildkite/test_areas/kernels.yaml

@@ -176,9 +176,9 @@ steps:
   timeout_in_minutes: 25
   device: h200_18gb
   source_file_dependencies:
-  - vllm/model_executor/layers/fla/ops/kda.py
-  - vllm/model_executor/layers/fla/ops/chunk_delta_h.py
-  - vllm/model_executor/layers/fla/ops/l2norm.py
+  - vllm/third_party/flash_linear_attention/ops/kda.py
+  - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py
+  - vllm/third_party/flash_linear_attention/ops/l2norm.py
   - tests/kernels/test_kda.py
   commands:
     - pytest -v -s kernels/test_kda.py

.buildkite/test_areas/lm_eval.yaml

@@ -103,7 +103,7 @@ steps:
   - vllm/transformers_utils/configs/qwen3_5_moe.py
   - vllm/model_executor/models/qwen3_next.py
   - vllm/model_executor/models/qwen3_next_mtp.py
-  - vllm/model_executor/layers/fla/ops/
+  - vllm/third_party/flash_linear_attention/ops/
   commands:
   - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt
 

.github/CODEOWNERS

@@ -172,7 +172,7 @@ mkdocs.yaml @hmellor
 # Kernels
 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep
 /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep
-/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy
+/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy
 
 # ROCm related: specify owner with write access to notify AMD folks for careful code review
 /vllm/**/*rocm* @tjtanaa @dllehr-amd

.yapfignore

@@ -1,2 +0,0 @@
-collect_env.py
-vllm/model_executor/layers/fla/ops/*.py

setup.py

@@ -1188,6 +1188,7 @@ package_data = {
         "entrypoints/serve/instrumentator/static/*.js",
         "entrypoints/serve/instrumentator/static/*.css",
         "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp",
+        "third_party/flash_linear_attention/LICENSE",
         # DeepGEMM JIT include headers (vendored via cmake)
         "third_party/deep_gemm/include/**/*.cuh",
         "third_party/deep_gemm/include/**/*.h",

tests/kernels/core/test_fused_rms_norm_gated.py

@@ -7,7 +7,7 @@ matching the eager triton kernel output."""
 import pytest
 import torch
 
-from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated
+from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated
 from vllm.utils.torch_utils import set_random_seed
 
 DTYPES = [torch.bfloat16]

tests/kernels/mamba/test_gdn_forward_core_split.py

@@ -53,11 +53,6 @@ from tests.v1.attention.utils import (  # noqa: E402
     create_vllm_config,
 )
 from vllm.config import set_current_vllm_config  # noqa: E402
-from vllm.model_executor.layers.fla.ops.index import (  # noqa: E402
-    prepare_chunk_indices,
-    prepare_chunk_offsets,
-)
-from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE  # noqa: E402
 from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn  # noqa: E402
 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (  # noqa: E402
     ChunkGatedDeltaRule,
@@ -66,6 +61,13 @@ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (  # noqa:
 from vllm.model_executor.layers.mamba.mamba_utils import (  # noqa: E402
     MambaStateShapeCalculator,
 )
+from vllm.third_party.flash_linear_attention.ops.index import (  # noqa: E402
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
+from vllm.third_party.flash_linear_attention.ops.utils import (  # noqa: E402
+    FLA_CHUNK_SIZE,
+)
 from vllm.v1.attention.backends.gdn_attn import (  # noqa: E402
     GDNAttentionMetadataBuilder,
 )

tests/kernels/mamba/test_gdn_prefill_cutedsl.py

@@ -17,17 +17,17 @@ if not (
         allow_module_level=True,
     )
 
-from vllm.model_executor.layers.fla.ops import (  # noqa: E402
-    chunk_gated_delta_rule,
-)
-from vllm.model_executor.layers.fla.ops.index import (  # noqa: E402
-    prepare_chunk_indices,
-    prepare_chunk_offsets,
-)
 from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import (  # noqa: E402
     chunk_gated_delta_rule_cutedsl,
     prepare_metadata_cutedsl,
 )
+from vllm.third_party.flash_linear_attention.ops import (  # noqa: E402
+    chunk_gated_delta_rule,
+)
+from vllm.third_party.flash_linear_attention.ops.index import (  # noqa: E402
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
 
 
 @pytest.mark.parametrize("num_seqs", [1, 5, 257])

tests/kernels/test_fla_layernorm_guard.py

@@ -5,7 +5,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from vllm.model_executor.layers.fla.ops.layernorm_guard import (
+from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
     layer_norm_fwd,
     layernorm_fn,
     rms_norm_ref,

tests/kernels/test_fused_gdn_post_conv.py

@@ -10,7 +10,7 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
+from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import (
     fused_post_conv_prep,
 )
 

tests/kernels/test_fused_recurrent_packed_decode.py

@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from vllm.model_executor.layers.fla.ops import (
+from vllm.third_party.flash_linear_attention.ops import (
     fused_recurrent_gated_delta_rule,
     fused_recurrent_gated_delta_rule_packed_decode,
 )

tests/kernels/test_fused_sigmoid_gating_delta_rule.py

@@ -5,11 +5,11 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from vllm.model_executor.layers.fla.ops import (
+from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops import (
     fused_recurrent_gated_delta_rule,
     fused_sigmoid_gating_delta_rule_update,
 )
-from vllm.platforms import current_platform
 from vllm.utils.torch_utils import set_random_seed
 
 DEVICE = current_platform.device_type

tests/kernels/test_kda.py

@@ -10,12 +10,12 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from vllm.model_executor.layers.fla.ops.kda import (
+from vllm.third_party.flash_linear_attention.ops.kda import (
     chunk_kda,
     chunk_kda_with_fused_gate,
     fused_kda_gate,
 )
-from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd
+from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd
 
 DEVICE = "cuda"
 

tools/pre_commit/mypy.py

@@ -128,7 +128,6 @@ EXCLUDE = [
     r"vllm/model_executor/models/[vV]",
     r"vllm/model_executor/models/[wW]",
     r"vllm/model_executor/models/[zZ]",
-    "vllm/model_executor/layers/fla/ops",
 ]
 
 

vllm/compilation/passes/fusion/matcher_utils.py

@@ -192,7 +192,7 @@ class MatcherRMSNormGated(MatcherCustomOp):
         z: torch.Tensor,
         weight: torch.Tensor,
     ) -> torch.Tensor:
-        from vllm.model_executor.layers.fla.ops.layernorm_guard import (
+        from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
             rmsnorm_fn,
         )
 

vllm/model_executor/layers/layernorm.py

@@ -286,7 +286,9 @@ class RMSNormGated(CustomOp):
     def forward_cuda(
         self, x: torch.Tensor, z: torch.Tensor | None = None
     ) -> torch.Tensor:
-        from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn
+        from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
+            rmsnorm_fn,
+        )
 
         return rmsnorm_fn(
             x,

vllm/model_executor/layers/mamba/gdn/kimi_gdn_linear_attn.py

@@ -15,16 +15,16 @@ from vllm.model_executor.custom_op import PluggableLayer
 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention
 from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader
 from vllm.model_executor.utils import set_weight_attrs
-from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
-from vllm.utils.torch_utils import direct_register_custom_op
-from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
-
-from ...fla.ops.kda import (
+from vllm.third_party.flash_linear_attention.ops.kda import (
     FusedRMSNormGated,
     chunk_kda_with_fused_gate,
     fused_kda_gate,
     fused_recurrent_kda,
 )
+from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig
+from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata
+
 from ...linear import (
     ColumnParallelLinear,
     ReplicatedLinear,

vllm/model_executor/layers/mamba/gdn/olmo_gdn_linear_attn.py

@@ -13,10 +13,6 @@ from vllm.distributed import (
 )
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.model_executor.custom_op import PluggableLayer
-from vllm.model_executor.layers.fla.ops import (
-    chunk_gated_delta_rule,
-    fused_recurrent_gated_delta_rule,
-)
 from vllm.model_executor.layers.layernorm import RMSNormGated
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -37,6 +33,10 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops import (
+    chunk_gated_delta_rule,
+    fused_recurrent_gated_delta_rule,
+)
 from vllm.triton_utils import tl, triton
 from vllm.triton_utils.allocation import set_triton_allocator
 from vllm.utils.torch_utils import direct_register_custom_op

vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py

@@ -20,16 +20,6 @@ from vllm.distributed import (
 from vllm.forward_context import ForwardContext, get_forward_context
 from vllm.logger import init_logger
 from vllm.model_executor.custom_op import CustomOp, PluggableLayer
-from vllm.model_executor.layers.fla.ops import (
-    chunk_gated_delta_rule as fla_chunk_gated_delta_rule,
-)
-from vllm.model_executor.layers.fla.ops import (
-    fused_post_conv_prep,
-    fused_recurrent_gated_delta_rule_packed_decode,
-    fused_sigmoid_gating_delta_rule_update,
-)
-from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd
-from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
 from vllm.model_executor.layers.layernorm import RMSNormGated
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -55,6 +45,16 @@ from vllm.model_executor.model_loader.weight_utils import (
 )
 from vllm.model_executor.utils import set_weight_attrs
 from vllm.platforms import current_platform
+from vllm.third_party.flash_linear_attention.ops import (
+    chunk_gated_delta_rule as fla_chunk_gated_delta_rule,
+)
+from vllm.third_party.flash_linear_attention.ops import (
+    fused_post_conv_prep,
+    fused_recurrent_gated_delta_rule_packed_decode,
+    fused_sigmoid_gating_delta_rule_update,
+)
+from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd
+from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE
 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
 from vllm.triton_utils import tl, triton
 from vllm.utils.torch_utils import (

vllm/model_executor/layers/mamba/linear/bailing_linear_attn.py

@@ -16,10 +16,6 @@ from vllm.distributed import (
 )
 from vllm.forward_context import get_forward_context
 from vllm.model_executor.custom_op import PluggableLayer
-from vllm.model_executor.layers.fla.ops.layernorm_guard import (
-    RMSNormGated,
-    layernorm_fn,
-)
 from vllm.model_executor.layers.layernorm import RMSNorm
 from vllm.model_executor.layers.linear import (
     ColumnParallelLinear,
@@ -33,6 +29,10 @@ from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (
     linear_attention_decode,
 )
 from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
+    RMSNormGated,
+    layernorm_fn,
+)
 from vllm.triton_utils import tl, triton
 from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata
 

vllm/model_executor/warmup/qwen_triton_warmup.py

@@ -274,7 +274,7 @@ def _warm_causal_conv1d_fwd_kernel(
 def _warm_fused_post_conv_kernel(
     device: torch.device, config: _QwenGDNWarmupConfig
 ) -> None:
-    from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (
+    from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import (  # noqa: E501
         fused_post_conv_prep,
     )
 
@@ -304,7 +304,7 @@ def _warm_fused_sigmoid_gating_delta_rule_update_kernel(
     device: torch.device,
     config: _QwenGDNWarmupConfig,
 ) -> None:
-    from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import (
+    from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import (
         fused_sigmoid_gating_delta_rule_update,
     )
 

vllm/third_party/flash_linear_attention/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

vllm/v1/attention/backends/gdn_attn.py

@@ -331,7 +331,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
         prefill_state_indices: torch.Tensor | None = None
         prefill_has_initial_state: torch.Tensor | None = None
         if num_prefills > 0:
-            from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE
+            from vllm.third_party.flash_linear_attention.ops.utils import (
+                FLA_CHUNK_SIZE,
+            )
 
             # In a mixed non-spec batch, decodes are peeled off to the recurrent
             # kernel (decode-first front slice), so build chunk metadata from the
@@ -371,7 +373,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
                 # Only prefill batches use FLA chunk ops.
                 # Pre-compute on CPU and async-copy to GPU to avoid
                 # GPU→CPU sync (.tolist()) in prepare_chunk_indices.
-                from vllm.model_executor.layers.fla.ops.index import (
+                from vllm.third_party.flash_linear_attention.ops.index import (
                     prepare_chunk_indices,
                     prepare_chunk_offsets,
                 )