626c90b2d5 · tduncan/Nvidia-Optimization-Stack-VLLM

.buildkite/test-amd.yaml

+5 -5

@@ -1444,9 +1444,9 @@ steps:

1444 optional: true

1445 working_dir: "/vllm-workspace/tests"

1446 source_file_dependencies:

1447 - - vllm/model_executor/layers/fla/ops/kda.py

1448 - - vllm/model_executor/layers/fla/ops/chunk_delta_h.py

1449 - - vllm/model_executor/layers/fla/ops/l2norm.py

1447 + - vllm/third_party/flash_linear_attention/ops/kda.py

1448 + - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py

1449 + - vllm/third_party/flash_linear_attention/ops/l2norm.py

1450 - tests/kernels/test_kda.py

1451 - vllm/platforms/rocm.py

1452 commands:

@@ -3234,7 +3234,7 @@ steps:

3234 - vllm/model_executor/models/qwen3.py

3235 - vllm/model_executor/models/qwen3_next.py

3236 - vllm/model_executor/models/qwen3_next_mtp.py

3237 - - vllm/model_executor/layers/fla/ops/

3237 + - vllm/third_party/flash_linear_attention/ops/

3238 - vllm/_aiter_ops.py

3239 - vllm/platforms/rocm.py

3240 commands:

@@ -3473,7 +3473,7 @@ steps:

3473 - vllm/model_executor/models/qwen3.py

3474 - vllm/model_executor/models/qwen3_next.py

3475 - vllm/model_executor/models/qwen3_next_mtp.py

3476 - - vllm/model_executor/layers/fla/ops/

3476 + - vllm/third_party/flash_linear_attention/ops/

3477 - vllm/_aiter_ops.py

3478 - vllm/v1/attention/backends/triton_attn.py

3479 - vllm/v1/attention/backends/rocm_attn.py

.buildkite/test_areas/kernels.yaml

+3 -3

@@ -176,9 +176,9 @@ steps:

176 timeout_in_minutes: 25

177 device: h200_18gb

178 source_file_dependencies:

179 - - vllm/model_executor/layers/fla/ops/kda.py

180 - - vllm/model_executor/layers/fla/ops/chunk_delta_h.py

181 - - vllm/model_executor/layers/fla/ops/l2norm.py

179 + - vllm/third_party/flash_linear_attention/ops/kda.py

180 + - vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py

181 + - vllm/third_party/flash_linear_attention/ops/l2norm.py

182 - tests/kernels/test_kda.py

183 commands:

184 - pytest -v -s kernels/test_kda.py

.buildkite/test_areas/lm_eval.yaml

+1 -1

@@ -103,7 +103,7 @@ steps:

103 - vllm/transformers_utils/configs/qwen3_5_moe.py

104 - vllm/model_executor/models/qwen3_next.py

105 - vllm/model_executor/models/qwen3_next_mtp.py

106 - - vllm/model_executor/layers/fla/ops/

106 + - vllm/third_party/flash_linear_attention/ops/

107 commands:

108 - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-qwen35-blackwell.txt

109

.github/CODEOWNERS

+1 -1

@@ -172,7 +172,7 @@ mkdocs.yaml @hmellor

172 # Kernels

173 /vllm/v1/attention/ops/chunked_prefill_paged_decode.py @tdoublep

174 /vllm/v1/attention/ops/triton_unified_attention.py @tdoublep

175 -/vllm/model_executor/layers/fla @ZJY0516 @vadiklyutiy

175 +/vllm/third_party/flash_linear_attention @ZJY0516 @vadiklyutiy

176

177 # ROCm related: specify owner with write access to notify AMD folks for careful code review

178 /vllm/**/*rocm* @tjtanaa @dllehr-amd

.yapfignore deleted

-2

@@ -1,2 +0,0 @@

1 -collect_env.py

2 -vllm/model_executor/layers/fla/ops/*.py

setup.py

+1

@@ -1188,6 +1188,7 @@ package_data = {

1188 "entrypoints/serve/instrumentator/static/*.js",

1189 "entrypoints/serve/instrumentator/static/*.css",

1190 "distributed/kv_transfer/kv_connector/v1/hf3fs/utils/*.cpp",

1191 + "third_party/flash_linear_attention/LICENSE",

1192 # DeepGEMM JIT include headers (vendored via cmake)

1193 "third_party/deep_gemm/include/**/*.cuh",

1194 "third_party/deep_gemm/include/**/*.h",

tests/kernels/core/test_fused_rms_norm_gated.py

+1 -1

@@ -7,7 +7,7 @@ matching the eager triton kernel output."""

7 import pytest

8 import torch

9

10 -from vllm.model_executor.layers.fla.ops.kda import FusedRMSNormGated

10 +from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated

11 from vllm.utils.torch_utils import set_random_seed

12

13 DTYPES = [torch.bfloat16]

tests/kernels/mamba/test_gdn_forward_core_split.py

+7 -5

@@ -53,11 +53,6 @@ from tests.v1.attention.utils import ( # noqa: E402

53 create_vllm_config,

54 )

55 from vllm.config import set_current_vllm_config # noqa: E402

56 -from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402

57 - prepare_chunk_indices,

58 - prepare_chunk_offsets,

59 -)

60 -from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE # noqa: E402

56 from vllm.model_executor.layers.mamba.gdn import qwen_gdn_linear_attn # noqa: E402

57 from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa: E402

58 ChunkGatedDeltaRule,

@@ -66,6 +61,13 @@ from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import ( # noqa:

61 from vllm.model_executor.layers.mamba.mamba_utils import ( # noqa: E402

62 MambaStateShapeCalculator,

63 )

64 +from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402

65 + prepare_chunk_indices,

66 + prepare_chunk_offsets,

67 +)

68 +from vllm.third_party.flash_linear_attention.ops.utils import ( # noqa: E402

69 + FLA_CHUNK_SIZE,

70 +)

71 from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402

72 GDNAttentionMetadataBuilder,

73 )

tests/kernels/mamba/test_gdn_prefill_cutedsl.py

+6 -6

@@ -17,17 +17,17 @@ if not (

17 allow_module_level=True,

18 )

19

20 -from vllm.model_executor.layers.fla.ops import ( # noqa: E402

20 +from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402

21 + chunk_gated_delta_rule_cutedsl,

22 + prepare_metadata_cutedsl,

23 +)

24 +from vllm.third_party.flash_linear_attention.ops import ( # noqa: E402

25 chunk_gated_delta_rule,

26 )

23 -from vllm.model_executor.layers.fla.ops.index import ( # noqa: E402

27 +from vllm.third_party.flash_linear_attention.ops.index import ( # noqa: E402

28 prepare_chunk_indices,

29 prepare_chunk_offsets,

30 )

27 -from vllm.model_executor.layers.mamba.ops.gdn_chunk_cutedsl import ( # noqa: E402

28 - chunk_gated_delta_rule_cutedsl,

29 - prepare_metadata_cutedsl,

30 -)

31

32

33 @pytest.mark.parametrize("num_seqs", [1, 5, 257])

tests/kernels/test_fla_layernorm_guard.py

+1 -1

@@ -5,7 +5,7 @@ import pytest

5 import torch

6 import torch.nn.functional as F

7

8 -from vllm.model_executor.layers.fla.ops.layernorm_guard import (

8 +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (

9 layer_norm_fwd,

10 layernorm_fn,

11 rms_norm_ref,

tests/kernels/test_fused_gdn_post_conv.py

+1 -1

@@ -10,7 +10,7 @@ import pytest

10 import torch

11 import torch.nn.functional as F

12

13 -from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (

13 +from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import (

14 fused_post_conv_prep,

15 )

16

tests/kernels/test_fused_recurrent_packed_decode.py

+1 -1

@@ -4,7 +4,7 @@

4 import pytest

5 import torch

6

7 -from vllm.model_executor.layers.fla.ops import (

7 +from vllm.third_party.flash_linear_attention.ops import (

8 fused_recurrent_gated_delta_rule,

9 fused_recurrent_gated_delta_rule_packed_decode,

10 )

tests/kernels/test_fused_sigmoid_gating_delta_rule.py

+2 -2

@@ -5,11 +5,11 @@ import pytest

5 import torch

6 import torch.nn.functional as F

7

8 -from vllm.model_executor.layers.fla.ops import (

8 +from vllm.platforms import current_platform

9 +from vllm.third_party.flash_linear_attention.ops import (

10 fused_recurrent_gated_delta_rule,

11 fused_sigmoid_gating_delta_rule_update,

12 )

12 -from vllm.platforms import current_platform

13 from vllm.utils.torch_utils import set_random_seed

14

15 DEVICE = current_platform.device_type

tests/kernels/test_kda.py

+2 -2

@@ -10,12 +10,12 @@ import pytest

10 import torch

11 import torch.nn.functional as F

12

13 -from vllm.model_executor.layers.fla.ops.kda import (

13 +from vllm.third_party.flash_linear_attention.ops.kda import (

14 chunk_kda,

15 chunk_kda_with_fused_gate,

16 fused_kda_gate,

17 )

18 -from vllm.model_executor.layers.fla.ops.l2norm import l2norm_fwd

18 +from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd

19

20 DEVICE = "cuda"

21

tools/pre_commit/mypy.py

-1

@@ -128,7 +128,6 @@ EXCLUDE = [

128 r"vllm/model_executor/models/[vV]",

129 r"vllm/model_executor/models/[wW]",

130 r"vllm/model_executor/models/[zZ]",

131 - "vllm/model_executor/layers/fla/ops",

131 ]

132

133

vllm/compilation/passes/fusion/matcher_utils.py

+1 -1

@@ -192,7 +192,7 @@ class MatcherRMSNormGated(MatcherCustomOp):

192 z: torch.Tensor,

193 weight: torch.Tensor,

194 ) -> torch.Tensor:

195 - from vllm.model_executor.layers.fla.ops.layernorm_guard import (

195 + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (

196 rmsnorm_fn,

197 )

198

vllm/model_executor/layers/layernorm.py

+3 -1

@@ -286,7 +286,9 @@ class RMSNormGated(CustomOp):

286 def forward_cuda(

287 self, x: torch.Tensor, z: torch.Tensor | None = None

288 ) -> torch.Tensor:

289 - from vllm.model_executor.layers.fla.ops.layernorm_guard import rmsnorm_fn

289 + from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (

290 + rmsnorm_fn,

291 + )

292

293 return rmsnorm_fn(

294 x,

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

+5 -5

@@ -15,16 +15,16 @@ from vllm.model_executor.custom_op import PluggableLayer

15 from vllm.model_executor.layers.mamba.gdn.base import GatedDeltaNetAttention

16 from vllm.model_executor.model_loader.weight_utils import sharded_weight_loader

17 from vllm.model_executor.utils import set_weight_attrs

18 -from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

19 -from vllm.utils.torch_utils import direct_register_custom_op

20 -from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

21 -

22 -from ...fla.ops.kda import (

18 +from vllm.third_party.flash_linear_attention.ops.kda import (

19 FusedRMSNormGated,

20 chunk_kda_with_fused_gate,

21 fused_kda_gate,

22 fused_recurrent_kda,

23 )

24 +from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig

25 +from vllm.utils.torch_utils import direct_register_custom_op

26 +from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadata

27 +

28 from ...linear import (

29 ColumnParallelLinear,

30 ReplicatedLinear,

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

+4 -4

@@ -13,10 +13,6 @@ from vllm.distributed import (

13 )

14 from vllm.forward_context import ForwardContext, get_forward_context

15 from vllm.model_executor.custom_op import PluggableLayer

16 -from vllm.model_executor.layers.fla.ops import (

17 - chunk_gated_delta_rule,

18 - fused_recurrent_gated_delta_rule,

19 -)

16 from vllm.model_executor.layers.layernorm import RMSNormGated

17 from vllm.model_executor.layers.linear import (

18 ColumnParallelLinear,

@@ -37,6 +33,10 @@ from vllm.model_executor.model_loader.weight_utils import (

33 )

34 from vllm.model_executor.utils import set_weight_attrs

35 from vllm.platforms import current_platform

36 +from vllm.third_party.flash_linear_attention.ops import (

37 + chunk_gated_delta_rule,

38 + fused_recurrent_gated_delta_rule,

39 +)

40 from vllm.triton_utils import tl, triton

41 from vllm.triton_utils.allocation import set_triton_allocator

42 from vllm.utils.torch_utils import direct_register_custom_op

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

+10 -10

@@ -20,16 +20,6 @@ from vllm.distributed import (

20 from vllm.forward_context import ForwardContext, get_forward_context

21 from vllm.logger import init_logger

22 from vllm.model_executor.custom_op import CustomOp, PluggableLayer

23 -from vllm.model_executor.layers.fla.ops import (

24 - chunk_gated_delta_rule as fla_chunk_gated_delta_rule,

25 -)

26 -from vllm.model_executor.layers.fla.ops import (

27 - fused_post_conv_prep,

28 - fused_recurrent_gated_delta_rule_packed_decode,

29 - fused_sigmoid_gating_delta_rule_update,

30 -)

31 -from vllm.model_executor.layers.fla.ops.chunk import l2norm_fwd

32 -from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE

23 from vllm.model_executor.layers.layernorm import RMSNormGated

24 from vllm.model_executor.layers.linear import (

25 ColumnParallelLinear,

@@ -55,6 +45,16 @@ from vllm.model_executor.model_loader.weight_utils import (

45 )

46 from vllm.model_executor.utils import set_weight_attrs

47 from vllm.platforms import current_platform

48 +from vllm.third_party.flash_linear_attention.ops import (

49 + chunk_gated_delta_rule as fla_chunk_gated_delta_rule,

50 +)

51 +from vllm.third_party.flash_linear_attention.ops import (

52 + fused_post_conv_prep,

53 + fused_recurrent_gated_delta_rule_packed_decode,

54 + fused_sigmoid_gating_delta_rule_update,

55 +)

56 +from vllm.third_party.flash_linear_attention.ops.chunk import l2norm_fwd

57 +from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE

58 from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig

59 from vllm.triton_utils import tl, triton

60 from vllm.utils.torch_utils import (

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

+4 -4

@@ -16,10 +16,6 @@ from vllm.distributed import (

16 )

17 from vllm.forward_context import get_forward_context

18 from vllm.model_executor.custom_op import PluggableLayer

19 -from vllm.model_executor.layers.fla.ops.layernorm_guard import (

20 - RMSNormGated,

21 - layernorm_fn,

22 -)

19 from vllm.model_executor.layers.layernorm import RMSNorm

20 from vllm.model_executor.layers.linear import (

21 ColumnParallelLinear,

@@ -33,6 +29,10 @@ from vllm.model_executor.layers.mamba.linear.minimax_linear_attn import (

29 linear_attention_decode,

30 )

31 from vllm.model_executor.layers.rotary_embedding import get_rope

32 +from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (

33 + RMSNormGated,

34 + layernorm_fn,

35 +)

36 from vllm.triton_utils import tl, triton

37 from vllm.v1.attention.backends.linear_attn import LinearAttentionMetadata

38

vllm/model_executor/warmup/qwen_triton_warmup.py

+2 -2

@@ -274,7 +274,7 @@ def _warm_causal_conv1d_fwd_kernel(

274 def _warm_fused_post_conv_kernel(

275 device: torch.device, config: _QwenGDNWarmupConfig

276 ) -> None:

277 - from vllm.model_executor.layers.fla.ops.fused_gdn_prefill_post_conv import (

277 + from vllm.third_party.flash_linear_attention.ops.fused_gdn_prefill_post_conv import ( # noqa: E501

278 fused_post_conv_prep,

279 )

280

@@ -304,7 +304,7 @@ def _warm_fused_sigmoid_gating_delta_rule_update_kernel(

304 device: torch.device,

305 config: _QwenGDNWarmupConfig,

306 ) -> None:

307 - from vllm.model_executor.layers.fla.ops.fused_sigmoid_gating import (

307 + from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import (

308 fused_sigmoid_gating_delta_rule_update,

309 )

310

vllm/third_party/flash_linear_attention/LICENSE new

+21

@@ -0,0 +1,21 @@

1 +MIT License

2 +

4 +

5 +Permission is hereby granted, free of charge, to any person obtaining a copy

6 +of this software and associated documentation files (the "Software"), to deal

7 +in the Software without restriction, including without limitation the rights

8 +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell

9 +copies of the Software, and to permit persons to whom the Software is

10 +furnished to do so, subject to the following conditions:

11 +

12 +The above copyright notice and this permission notice shall be included in all

13 +copies or substantial portions of the Software.

14 +

15 +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR

16 +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,

17 +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE

18 +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER

19 +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,

20 +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE

21 +SOFTWARE.

vllm/third_party/flash_linear_attention/__init__.py renamed

vllm/third_party/flash_linear_attention/ops/__init__.py renamed

vllm/third_party/flash_linear_attention/ops/chunk.py renamed

vllm/third_party/flash_linear_attention/ops/chunk_delta_h.py renamed

vllm/third_party/flash_linear_attention/ops/chunk_o.py renamed

vllm/third_party/flash_linear_attention/ops/chunk_scaled_dot_kkt.py renamed

vllm/third_party/flash_linear_attention/ops/cumsum.py renamed

vllm/third_party/flash_linear_attention/ops/fused_gdn_prefill_post_conv.py renamed

vllm/third_party/flash_linear_attention/ops/fused_recurrent.py renamed

vllm/third_party/flash_linear_attention/ops/fused_sigmoid_gating.py renamed

vllm/third_party/flash_linear_attention/ops/index.py renamed

vllm/third_party/flash_linear_attention/ops/kda.py renamed

vllm/third_party/flash_linear_attention/ops/l2norm.py renamed

vllm/third_party/flash_linear_attention/ops/layernorm_guard.py renamed

vllm/third_party/flash_linear_attention/ops/op.py renamed

vllm/third_party/flash_linear_attention/ops/solve_tril.py renamed

vllm/third_party/flash_linear_attention/ops/utils.py renamed

vllm/third_party/flash_linear_attention/ops/wy_fast.py renamed

vllm/v1/attention/backends/gdn_attn.py

+4 -2

@@ -331,7 +331,9 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

331 prefill_state_indices: torch.Tensor | None = None

332 prefill_has_initial_state: torch.Tensor | None = None

333 if num_prefills > 0:

334 - from vllm.model_executor.layers.fla.ops.utils import FLA_CHUNK_SIZE

334 + from vllm.third_party.flash_linear_attention.ops.utils import (

335 + FLA_CHUNK_SIZE,

336 + )

337

338 # In a mixed non-spec batch, decodes are peeled off to the recurrent

339 # kernel (decode-first front slice), so build chunk metadata from the

@@ -371,7 +373,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]

373 # Only prefill batches use FLA chunk ops.

374 # Pre-compute on CPU and async-copy to GPU to avoid

375 # GPU→CPU sync (.tolist()) in prepare_chunk_indices.

374 - from vllm.model_executor.layers.fla.ops.index import (

376 + from vllm.third_party.flash_linear_attention.ops.index import (

377 prepare_chunk_indices,

378 prepare_chunk_offsets,

379 )

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