A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 13:33 UTC
[5/N][Attention] Finish eliminating `vllm/attention` folder (#32064)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Jan 27, 2026 at 10:02 UTC
a608b4c6c2b9098f7500ed305f47cbe1375df5ee
151 files changed
+585
-527
.buildkite/test-amd.yaml
+2
-1
@@ -640,8 +640,9 @@ steps:
640
# grade: Blocking
641
source_file_dependencies:
642
- csrc/attention/
643
- - vllm/attention
643
- vllm/v1/attention
644
+ # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
645
+ - vllm/model_executor/layers/attention
646
- tests/kernels/attention
647
commands:
648
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
.buildkite/test-pipeline.yaml
+2
-1
@@ -568,8 +568,9 @@ steps:
568
mirror_hardwares: [amdexperimental]
569
source_file_dependencies:
570
- csrc/attention/
571
- - vllm/attention
571
- vllm/v1/attention
572
+ # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
573
+ - vllm/model_executor/layers/attention
574
- tests/kernels/attention
575
commands:
576
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
.buildkite/test_areas/kernels.yaml
+2
-1
@@ -15,8 +15,9 @@ steps:
15
timeout_in_minutes: 35
16
source_file_dependencies:
17
- csrc/attention/
18
- - vllm/attention
18
- vllm/v1/attention
19
+ # TODO: remove this dependency (https://github.com/vllm-project/vllm/issues/32267)
20
+ - vllm/model_executor/layers/attention
21
- tests/kernels/attention
22
commands:
23
- pytest -v -s kernels/attention --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
.github/CODEOWNERS
+1
-1
@@ -2,8 +2,8 @@
2
# for more info about CODEOWNERS file
3
4
# This lists cover the "core" components of vLLM that require careful review
5
-/vllm/attention @LucasWilkinson
5
/vllm/executor/executor_base.py @zhuohan123 @youkaichao @alexm-redhat @njhill @22quinn
6
+/vllm/model_executor/layers/attention @LucasWilkinson
7
/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety
8
/vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety
9
/vllm/model_executor/layers/mamba @tdoublep
docs/contributing/model/basic.md
+1
-1
@@ -29,7 +29,7 @@ The initialization code should look like this:
29
```python
30
from torch import nn
31
from vllm.config import VllmConfig
32
- from vllm.attention.layer import Attention
32
+ from vllm.model_executor.layers.attention import Attention
33
34
class MyAttention(nn.Module):
35
def __init__(self, vllm_config: VllmConfig, prefix: str):
docs/design/custom_op.md
+1
-1
@@ -271,7 +271,7 @@ Taking `MMEncoderAttention` as an example:
271
??? code
272
273
```python
274
- from vllm.attention.layers.mm_encoder_attention import MMEncoderAttention
274
+ from vllm.model_executor.layers.attention import MMEncoderAttention
275
from vllm.model_executor.custom_op import CustomOp
276
277
tests/compile/test_fusion_attn.py
+1
-1
@@ -21,7 +21,6 @@ from tests.compile.fusion_test_utils import (
21
from tests.utils import flat_product
22
from tests.v1.attention.utils import BatchSpec, create_common_attn_metadata
23
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
24
-from vllm.attention.layer import Attention
24
from vllm.compilation.fusion_attn import ATTN_OP, AttnFusionPass
25
from vllm.compilation.fx_utils import find_op_nodes
26
from vllm.compilation.matcher_utils import QUANT_OPS
@@ -40,6 +39,7 @@ from vllm.config import (
39
set_current_vllm_config,
40
)
41
from vllm.forward_context import get_forward_context, set_forward_context
42
+from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.quantization.utils.quant_utils import (
44
QuantKey,
45
kFp8StaticTensorSym,
tests/compile/test_qk_norm_rope_fusion.py
+1
-1
@@ -5,7 +5,6 @@ import pytest
5
import torch
6
7
from tests.compile.backend import TestBackend
8
-from vllm.attention.layer import Attention
8
from vllm.compilation.matcher_utils import FLASHINFER_ROTARY_OP, RMS_OP, ROTARY_OP
9
from vllm.compilation.noop_elimination import NoOpEliminationPass
10
from vllm.compilation.post_cleanup import PostCleanupPass
@@ -21,6 +20,7 @@ from vllm.config import (
20
VllmConfig,
21
set_current_vllm_config,
22
)
23
+from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
26
from vllm.platforms import current_platform
tests/kernels/attention/test_attention.py
+1
-2
@@ -9,8 +9,7 @@ import torch
9
from tests.kernels.allclose_default import get_default_atol, get_default_rtol
10
from tests.kernels.utils import opcheck
11
from vllm import _custom_ops as ops
12
-from vllm.attention.layer import Attention
13
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
12
+from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
13
from vllm.platforms import current_platform
14
from vllm.utils.mem_utils import get_max_shared_memory_bytes
15
from vllm.utils.torch_utils import set_random_seed
tests/kernels/attention/test_mha_attn.py
+1
-1
@@ -12,7 +12,7 @@ from unittest.mock import patch
12
import pytest
13
import torch
14
15
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
15
+from vllm.model_executor.layers.attention import MMEncoderAttention
16
from vllm.platforms import current_platform
17
from vllm.platforms.cpu import CpuPlatform
18
from vllm.platforms.cuda import CudaPlatform
tests/v1/worker/test_gpu_model_runner.py
+1
-1
@@ -5,7 +5,6 @@ import numpy as np
5
import pytest
6
import torch
7
8
-from vllm.attention.layer import Attention
8
from vllm.config import (
9
AttentionConfig,
10
CacheConfig,
@@ -19,6 +18,7 @@ from vllm.distributed.parallel_state import (
18
init_distributed_environment,
19
initialize_model_parallel,
20
)
21
+from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
23
from vllm.platforms import current_platform
24
from vllm.sampling_params import SamplingParams
tests/v1/worker/test_utils.py
+3
-3
@@ -7,7 +7,7 @@ from vllm.v1.worker.utils import bind_kv_cache
7
8
9
def test_bind_kv_cache(default_vllm_config):
10
- from vllm.attention.layer import Attention
10
+ from vllm.model_executor.layers.attention import Attention
11
12
ctx = {
13
"layers.0.self_attn": Attention(32, 128, 0.1, prefix="layers.0.self_attn"),
@@ -35,7 +35,7 @@ def test_bind_kv_cache(default_vllm_config):
35
36
37
def test_bind_kv_cache_non_attention(default_vllm_config):
38
- from vllm.attention.layer import Attention
38
+ from vllm.model_executor.layers.attention import Attention
39
40
# example from Jamba PP=2
41
ctx = {
@@ -58,7 +58,7 @@ def test_bind_kv_cache_non_attention(default_vllm_config):
58
59
60
def test_bind_kv_cache_draft_model(default_vllm_config):
61
- from vllm.attention.layer import Attention
61
+ from vllm.model_executor.layers.attention import Attention
62
63
layer_names = [
64
"model.layers.0.attn",
tools/pre_commit/mypy.py
-1
@@ -58,7 +58,6 @@ FILES = [
58
SEPARATE_GROUPS = [
59
"tests",
60
# v0 related
61
- "vllm/attention",
61
"vllm/compilation",
62
"vllm/lora",
63
"vllm/model_executor",
vllm/attention/__init__.py
vllm/attention/utils/__init__.py
vllm/attention/utils/kv_sharing_utils.py
deleted
-33
@@ -1,33 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-def validate_kv_sharing_target(
4
- current_layer_name, target_layer_name, static_forward_context
5
-):
6
- error_msg = (
7
- f"Specified KV sharing target layer for {current_layer_name} "
8
- f"is not valid: target layer {target_layer_name} "
9
- )
10
-
11
- if current_layer_name == target_layer_name:
12
- raise ValueError(error_msg + "cannot be the same as the current layer.")
13
-
14
- if target_layer_name not in static_forward_context:
15
- from vllm.model_executor.models.utils import extract_layer_index
16
-
17
- # If target layer name is not in the static fwd context, it means either
18
- # a) the target layer does not come BEFORE the current layer, or
19
- # b) the target layer is not an Attention layer that exists in the model
20
- current_layer_idx = extract_layer_index(current_layer_name)
21
- target_layer_idx = extract_layer_index(target_layer_name)
22
- if current_layer_idx <= target_layer_idx:
23
- raise ValueError(error_msg + "must come before the current layer.")
24
- else:
25
- raise ValueError(error_msg + "is not a valid Attention layer in the model.")
26
-
27
- # Currently KV sharing is only supported between layers of the same type
28
- target_layer_attn_type = static_forward_context[target_layer_name].attn_type
29
- expected = static_forward_context[current_layer_name].attn_type
30
- if target_layer_attn_type != expected:
31
- raise ValueError(
32
- error_msg + f"must be the same type as the current layer ({expected})."
33
- )
vllm/compilation/fusion_attn.py
+1
-1
@@ -11,9 +11,9 @@ from torch import fx
11
from torch._higher_order_ops.auto_functionalize import auto_functionalized
12
from torch._inductor.pattern_matcher import PatternMatcherPass
13
14
-from vllm.attention.layer import Attention
14
from vllm.config import VllmConfig, get_layers_from_vllm_config
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
QuantKey,
19
kNvfp4Dynamic,
vllm/compilation/qk_norm_rope_fusion.py
+1
-1
@@ -10,9 +10,9 @@ from torch import fx
10
from torch._higher_order_ops.auto_functionalize import auto_functionalized
11
from torch._inductor.pattern_matcher import PatternMatcherPass
12
13
-from vllm.attention.layer import Attention
13
from vllm.config import VllmConfig, get_layers_from_vllm_config
14
from vllm.logger import init_logger
15
+from vllm.model_executor.layers.attention import Attention
16
from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
17
18
from .fusion import empty_bf16, empty_fp32, empty_i64
vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py
+1
-1
@@ -8,7 +8,6 @@ from typing import Any
8
9
import torch
10
11
-from vllm.attention.layer import Attention
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
13
from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
@@ -25,6 +24,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import (
24
)
25
from vllm.forward_context import ForwardContext
26
from vllm.logger import init_logger
27
+from vllm.model_executor.layers.attention import Attention
28
from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata
29
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
30
from vllm.v1.core.kv_cache_utils import BlockHash
vllm/model_executor/layers/attention/__init__.py
+26
@@ -0,0 +1,26 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.model_executor.layers.attention.attention import Attention
5
+from vllm.model_executor.layers.attention.chunked_local_attention import (
6
+ ChunkedLocalAttention,
7
+)
8
+from vllm.model_executor.layers.attention.cross_attention import CrossAttention
9
+from vllm.model_executor.layers.attention.encoder_only_attention import (
10
+ EncoderOnlyAttention,
11
+)
12
+from vllm.model_executor.layers.attention.mla_attention import MLAAttention
13
+from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
14
+from vllm.model_executor.layers.attention.static_sink_attention import (
15
+ StaticSinkAttention,
16
+)
17
+
18
+__all__ = [
19
+ "Attention",
20
+ "ChunkedLocalAttention",
21
+ "CrossAttention",
22
+ "EncoderOnlyAttention",
23
+ "MLAAttention",
24
+ "MMEncoderAttention",
25
+ "StaticSinkAttention",
26
+]
vllm/model_executor/layers/attention/attention.py
renamed
+42
-315
@@ -1,23 +1,22 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Attention layer."""
3
5
-from typing import cast
4
+from typing import TYPE_CHECKING
5
6
import torch
7
import torch.nn as nn
8
9
import vllm.envs as envs
11
-from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
12
-from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
10
from vllm.config import CacheConfig, get_current_vllm_config
11
from vllm.config.vllm import VllmConfig
12
from vllm.forward_context import ForwardContext, get_forward_context
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.attention.kv_transfer_utils import (
15
+ maybe_transfer_kv_layer,
16
+)
17
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
18
from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
19
from vllm.model_executor.layers.linear import (
20
- ColumnParallelLinear,
20
UnquantizedLinearMethod,
21
)
22
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -33,20 +32,54 @@ from vllm.utils.torch_utils import (
32
from vllm.v1.attention.backend import (
33
AttentionBackend,
34
AttentionType,
36
- MLAAttentionImpl,
35
)
36
from vllm.v1.attention.backends.registry import AttentionBackendEnum
37
from vllm.v1.attention.selector import get_attn_backend
38
from vllm.v1.kv_cache_interface import (
39
FullAttentionSpec,
40
KVCacheSpec,
43
- MLAAttentionSpec,
41
SlidingWindowSpec,
42
)
43
44
+if TYPE_CHECKING:
45
+ from vllm.model_executor.layers.attention import MLAAttention
46
+
47
logger = init_logger(__name__)
48
49
50
+def validate_kv_sharing_target(
51
+ current_layer_name, target_layer_name, static_forward_context
52
+):
53
+ error_msg = (
54
+ f"Specified KV sharing target layer for {current_layer_name} "
55
+ f"is not valid: target layer {target_layer_name} "
56
+ )
57
+
58
+ if current_layer_name == target_layer_name:
59
+ raise ValueError(error_msg + "cannot be the same as the current layer.")
60
+
61
+ if target_layer_name not in static_forward_context:
62
+ from vllm.model_executor.models.utils import extract_layer_index
63
+
64
+ # If target layer name is not in the static fwd context, it means either
65
+ # a) the target layer does not come BEFORE the current layer, or
66
+ # b) the target layer is not an Attention layer that exists in the model
67
+ current_layer_idx = extract_layer_index(current_layer_name)
68
+ target_layer_idx = extract_layer_index(target_layer_name)
69
+ if current_layer_idx <= target_layer_idx:
70
+ raise ValueError(error_msg + "must come before the current layer.")
71
+ else:
72
+ raise ValueError(error_msg + "is not a valid Attention layer in the model.")
73
+
74
+ # Currently KV sharing is only supported between layers of the same type
75
+ target_layer_attn_type = static_forward_context[target_layer_name].attn_type
76
+ expected = static_forward_context[current_layer_name].attn_type
77
+ if target_layer_attn_type != expected:
78
+ raise ValueError(
79
+ error_msg + f"must be the same type as the current layer ({expected})."
80
+ )
81
+
82
+
83
def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
84
"""Returns whether the quantization method should load quantized weights."""
85
return quant_method is not None and not isinstance(
@@ -493,236 +526,6 @@ class Attention(nn.Module, AttentionLayerBase):
526
)
527
528
496
-class MLAAttention(nn.Module, AttentionLayerBase):
497
- """Multi-Head Latent Attention layer.
498
-
499
- This class takes query, and compressed key/value tensors as input.
500
- The class does the following:
501
-
502
- 1. Store the input key and value tensors in the KV cache.
503
- 2. Perform (multi-head/multi-query/grouped-query) attention.
504
- 3. Return the output tensor.
505
- """
506
-
507
- def __init__(
508
- self,
509
- num_heads: int,
510
- scale: float,
511
- qk_nope_head_dim: int,
512
- qk_rope_head_dim: int,
513
- v_head_dim: int,
514
- q_lora_rank: int | None,
515
- kv_lora_rank: int,
516
- kv_b_proj: ColumnParallelLinear,
517
- cache_config: CacheConfig | None = None,
518
- quant_config: QuantizationConfig | None = None,
519
- prefix: str = "",
520
- use_sparse: bool = False,
521
- indexer: object | None = None,
522
- **extra_impl_args,
523
- ):
524
- super().__init__()
525
- self.num_heads = num_heads
526
- self.scale = scale
527
- self.qk_nope_head_dim = qk_nope_head_dim
528
- self.qk_rope_head_dim = qk_rope_head_dim
529
- self.v_head_dim = v_head_dim
530
- self.q_lora_rank = q_lora_rank
531
- self.kv_lora_rank = kv_lora_rank
532
- self.head_size = kv_lora_rank + qk_rope_head_dim
533
- self.layer_name = prefix
534
-
535
- if cache_config is not None:
536
- kv_cache_dtype = cache_config.cache_dtype
537
- block_size = cache_config.block_size
538
- calculate_kv_scales = cache_config.calculate_kv_scales
539
- else:
540
- kv_cache_dtype = "auto"
541
- block_size = 16
542
- calculate_kv_scales = False
543
- self.quant_config = quant_config
544
-
545
- # Initialize KV cache quantization attributes
546
- self.kv_cache_dtype = kv_cache_dtype
547
- self.calculate_kv_scales = calculate_kv_scales
548
- _init_kv_cache_quant(self, quant_config, prefix)
549
-
550
- dtype = torch.get_default_dtype()
551
- self.attn_backend = get_attn_backend(
552
- self.head_size,
553
- dtype,
554
- kv_cache_dtype,
555
- block_size,
556
- use_mla=True,
557
- use_sparse=use_sparse,
558
- )
559
-
560
- if (
561
- cache_config is not None
562
- and cache_config.enable_prefix_caching
563
- and vllm_is_batch_invariant()
564
- and (
565
- self.attn_backend.get_name() == "TRITON_MLA"
566
- or self.attn_backend.get_name() == "FLASHINFER"
567
- )
568
- ):
569
- logger.warning_once(
570
- "Disabling prefix caching for TRITON_MLA / FLASHINFER "
571
- "with batch invariance, as it is not yet supported.",
572
- scope="local",
573
- )
574
- cache_config.enable_prefix_caching = False
575
-
576
- impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
577
- self.impl = impl_cls(
578
- num_heads=self.num_heads,
579
- head_size=self.head_size,
580
- scale=self.scale,
581
- num_kv_heads=1,
582
- alibi_slopes=None,
583
- sliding_window=None,
584
- kv_cache_dtype=self.kv_cache_dtype,
585
- logits_soft_cap=None,
586
- attn_type=AttentionType.DECODER,
587
- kv_sharing_target_layer_name=None,
588
- # MLA Args
589
- q_lora_rank=self.q_lora_rank,
590
- kv_lora_rank=self.kv_lora_rank,
591
- qk_nope_head_dim=self.qk_nope_head_dim,
592
- qk_rope_head_dim=self.qk_rope_head_dim,
593
- qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
594
- v_head_dim=self.v_head_dim,
595
- kv_b_proj=kv_b_proj,
596
- indexer=indexer,
597
- **extra_impl_args,
598
- )
599
-
600
- self.use_direct_call = not current_platform.opaque_attention_op()
601
-
602
- compilation_config = get_current_vllm_config().compilation_config
603
- if prefix in compilation_config.static_forward_context:
604
- raise ValueError(f"Duplicate layer name: {prefix}")
605
- compilation_config.static_forward_context[prefix] = self
606
-
607
- self.kv_cache = [
608
- torch.tensor([])
609
- for _ in range(
610
- get_current_vllm_config().parallel_config.pipeline_parallel_size
611
- )
612
- ]
613
-
614
- self.use_sparse = use_sparse
615
-
616
- # Initialize q/k/v range constants.
617
- self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
618
- self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
619
- self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
620
-
621
- def forward(
622
- self,
623
- q: torch.Tensor,
624
- kv_c_normed: torch.Tensor,
625
- k_pe: torch.Tensor,
626
- output_shape: torch.Size | None = None,
627
- ) -> torch.Tensor:
628
- if self.calculate_kv_scales:
629
- torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name)
630
-
631
- if self.use_direct_call:
632
- forward_context: ForwardContext = get_forward_context()
633
- attn_metadata = forward_context.attn_metadata
634
- if isinstance(attn_metadata, dict):
635
- attn_metadata = attn_metadata[self.layer_name]
636
- self_kv_cache = self.kv_cache[forward_context.virtual_engine]
637
-
638
- if self.attn_backend.accept_output_buffer:
639
- output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
640
- self.impl.forward(
641
- self,
642
- q,
643
- kv_c_normed,
644
- k_pe,
645
- self_kv_cache,
646
- attn_metadata,
647
- output=output,
648
- )
649
- return output
650
- else:
651
- return self.impl.forward(
652
- self, q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
653
- )
654
- else:
655
- if self.attn_backend.accept_output_buffer:
656
- output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
657
- torch.ops.vllm.unified_mla_attention_with_output(
658
- q,
659
- kv_c_normed,
660
- k_pe,
661
- output,
662
- self.layer_name,
663
- )
664
- return output
665
- else:
666
- return torch.ops.vllm.unified_mla_attention(
667
- q,
668
- kv_c_normed,
669
- k_pe,
670
- self.layer_name,
671
- )
672
-
673
- def process_weights_after_loading(self, act_dtype: torch.dtype):
674
- if hasattr(self.impl, "process_weights_after_loading"):
675
- self.impl.process_weights_after_loading(act_dtype)
676
-
677
- # If we should not load quant weights, we initialize the scales to 1.0
678
- # as the default value. See [Note: Register q/k/v/prob scales in state dict]
679
- # for more details.
680
- quant_method = (
681
- self.quant_config.get_quant_method(self, prefix=self.layer_name)
682
- if self.quant_config
683
- else None
684
- )
685
- if not should_load_quant_weights(quant_method):
686
- set_default_quant_scales(self, register_buffer=False)
687
-
688
- def calc_kv_scales(
689
- self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
690
- ) -> None:
691
- """Optional scale calculation for MLA inputs.
692
-
693
- Mirrors Attention.calc_kv_scales. Not all MLA backends require this
694
- """
695
- # Use safe defaults if ranges are not present
696
- q_range = getattr(self, "q_range", torch.tensor(1.0))
697
- k_range = getattr(self, "k_range", torch.tensor(1.0))
698
- v_range = getattr(self, "v_range", torch.tensor(1.0))
699
-
700
- self._q_scale.copy_(torch.abs(q).max() / q_range)
701
- # kv_c_normed is the compressed KV representation; use it for k/v
702
- kv_abs_max = torch.abs(kv_c_normed).max()
703
- self._k_scale.copy_(kv_abs_max / k_range)
704
- self._v_scale.copy_(kv_abs_max / v_range)
705
- self._q_scale_float = self._q_scale.item()
706
- self._k_scale_float = self._k_scale.item()
707
- self._v_scale_float = self._v_scale.item()
708
- self.calculate_kv_scales = False
709
-
710
- def get_attn_backend(self) -> type[AttentionBackend]:
711
- return self.attn_backend
712
-
713
- def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
714
- kv_cache_dtype = kv_cache_dtype_str_to_dtype(
715
- self.kv_cache_dtype, vllm_config.model_config
716
- )
717
- return MLAAttentionSpec(
718
- block_size=vllm_config.cache_config.block_size,
719
- num_kv_heads=1,
720
- head_size=self.head_size,
721
- dtype=kv_cache_dtype,
722
- cache_dtype_str=vllm_config.cache_config.cache_dtype,
723
- )
724
-
725
-
529
def maybe_calc_kv_scales(
530
query: torch.Tensor,
531
key: torch.Tensor,
@@ -759,7 +562,7 @@ direct_register_custom_op(
562
563
def get_attention_context(
564
layer_name: str,
762
-) -> tuple[dict | object | None, Attention | MLAAttention, torch.Tensor]:
565
+) -> tuple[dict | object | None, "Attention | MLAAttention", torch.Tensor]:
566
"""Extract attention context for a given layer.
567
568
This helper function extracts the attention metadata, attention layer
@@ -782,7 +585,7 @@ def get_attention_context(
585
attn_metadata = forward_context.attn_metadata
586
if isinstance(attn_metadata, dict):
587
attn_metadata = attn_metadata[layer_name]
785
- attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
588
+ attn_layer = forward_context.no_compile_layers[layer_name]
589
kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
590
return attn_metadata, attn_layer, kv_cache
591
@@ -914,79 +717,3 @@ direct_register_custom_op(
717
mutates_args=["output", "output_block_scale"],
718
fake_impl=unified_attention_with_output_fake,
719
)
917
-
918
-
919
-@maybe_transfer_kv_layer
920
-def unified_mla_attention(
921
- q: torch.Tensor,
922
- kv_c_normed: torch.Tensor,
923
- k_pe: torch.Tensor,
924
- layer_name: str,
925
-) -> torch.Tensor:
926
- attn_metadata, self, kv_cache = get_attention_context(layer_name)
927
- output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)
928
-
929
- return output
930
-
931
-
932
-def unified_mla_attention_fake(
933
- q: torch.Tensor,
934
- kv_c_normed: torch.Tensor,
935
- k_pe: torch.Tensor,
936
- layer_name: str,
937
-) -> torch.Tensor:
938
- return torch.empty_like(q).contiguous()
939
-
940
-
941
-direct_register_custom_op(
942
- op_name="unified_mla_attention",
943
- op_func=unified_mla_attention,
944
- mutates_args=[],
945
- fake_impl=unified_mla_attention_fake,
946
- dispatch_key=current_platform.dispatch_key,
947
-)
948
-
949
-
950
-@maybe_transfer_kv_layer
951
-def unified_mla_attention_with_output(
952
- q: torch.Tensor,
953
- kv_c_normed: torch.Tensor,
954
- k_pe: torch.Tensor,
955
- output: torch.Tensor,
956
- layer_name: str,
957
- output_scale: torch.Tensor | None = None,
958
- output_block_scale: torch.Tensor | None = None,
959
-) -> None:
960
- attn_metadata, self, kv_cache = get_attention_context(layer_name)
961
- self.impl.forward(
962
- self,
963
- q,
964
- kv_c_normed,
965
- k_pe,
966
- kv_cache,
967
- attn_metadata,
968
- output=output,
969
- output_scale=output_scale,
970
- output_block_scale=output_block_scale,
971
- )
972
-
973
-
974
-def unified_mla_attention_with_output_fake(
975
- q: torch.Tensor,
976
- kv_c_normed: torch.Tensor,
977
- k_pe: torch.Tensor,
978
- output: torch.Tensor,
979
- layer_name: str,
980
- output_scale: torch.Tensor | None = None,
981
- output_block_scale: torch.Tensor | None = None,
982
-) -> None:
983
- return
984
-
985
-
986
-direct_register_custom_op(
987
- op_name="unified_mla_attention_with_output",
988
- op_func=unified_mla_attention_with_output,
989
- mutates_args=["output", "output_block_scale"],
990
- fake_impl=unified_mla_attention_with_output_fake,
991
- dispatch_key=current_platform.dispatch_key,
992
-)
vllm/model_executor/layers/attention/chunked_local_attention.py
+1
-1
@@ -4,9 +4,9 @@ import functools
4
5
import torch
6
7
-from vllm.attention.layer import Attention
7
from vllm.config import CacheConfig
8
from vllm.config.vllm import VllmConfig
9
+from vllm.model_executor.layers.attention import Attention
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
from vllm.v1.attention.backend import (
12
AttentionBackend,
vllm/model_executor/layers/attention/cross_attention.py
+1
-1
@@ -6,9 +6,9 @@ from copy import copy
6
import numpy as np
7
import torch
8
9
-from vllm.attention.layer import Attention
9
from vllm.config import CacheConfig, VllmConfig
10
from vllm.logger import init_logger
11
+from vllm.model_executor.layers.attention import Attention
12
from vllm.utils.math_utils import cdiv
13
from vllm.v1.attention.backend import (
14
AttentionBackend,
vllm/model_executor/layers/attention/encoder_only_attention.py
+1
-1
@@ -5,9 +5,9 @@ from copy import copy
5
6
import torch
7
8
-from vllm.attention.layer import Attention
8
from vllm.config import CacheConfig
9
from vllm.config.vllm import VllmConfig
10
+from vllm.model_executor.layers.attention import Attention
11
from vllm.v1.attention.backend import (
12
AttentionBackend,
13
AttentionMetadata,
vllm/model_executor/layers/attention/kv_transfer_utils.py
renamed
+1
-1
@@ -19,7 +19,7 @@ def maybe_transfer_kv_layer(func: Callable) -> Callable:
19
On exit: saves the KV layer to the connector.
20
"""
21
# Import at runtime to avoid circular dependency
22
- from vllm.attention.layer import get_attention_context
22
+ from vllm.model_executor.layers.attention.attention import get_attention_context
23
24
# Inspect the signature ONCE when the decorator is applied.
25
sig = inspect.signature(func)
vllm/model_executor/layers/attention/mla_attention.py
+354
-17
@@ -191,24 +191,38 @@ import functools
191
from abc import abstractmethod
192
from dataclasses import dataclass, field
193
from enum import Enum
194
-from typing import ClassVar, Generic, TypeVar
194
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast
195
+
196
+if TYPE_CHECKING:
197
+ from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
198
199
import torch
200
+import torch.nn as nn
201
from tqdm import tqdm
202
203
+import vllm.envs as envs
204
from vllm import _custom_ops as ops
200
-from vllm import envs
205
from vllm._aiter_ops import rocm_aiter_ops
202
-from vllm.config import ModelConfig, VllmConfig, get_current_vllm_config
206
+from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
207
from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
208
+from vllm.forward_context import ForwardContext, get_forward_context
209
from vllm.logger import init_logger
210
from vllm.model_executor.custom_op import CustomOp
206
-from vllm.model_executor.layers.batch_invariant import (
207
- vllm_is_batch_invariant,
211
+from vllm.model_executor.layers.attention.attention import (
212
+ _init_kv_cache_quant,
213
+ get_attention_context,
214
+ set_default_quant_scales,
215
+ should_load_quant_weights,
216
+)
217
+from vllm.model_executor.layers.attention.kv_transfer_utils import (
218
+ maybe_transfer_kv_layer,
219
)
220
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
221
+from vllm.model_executor.layers.batch_invariant import vllm_is_batch_invariant
222
from vllm.model_executor.layers.linear import (
223
ColumnParallelLinear,
224
)
225
+from vllm.model_executor.layers.quantization import QuantizationConfig
226
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
227
from vllm.model_executor.layers.quantization.utils.quant_utils import (
228
GroupShape,
@@ -217,11 +231,16 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
231
from vllm.platforms import current_platform
232
from vllm.utils.flashinfer import has_nvidia_artifactory
233
from vllm.utils.math_utils import cdiv, round_down
234
+from vllm.utils.torch_utils import (
235
+ direct_register_custom_op,
236
+ kv_cache_dtype_str_to_dtype,
237
+)
238
from vllm.v1.attention.backend import (
239
AttentionBackend,
240
AttentionLayer,
241
AttentionMetadata,
242
AttentionMetadataBuilder,
243
+ AttentionType,
244
CommonAttentionMetadata,
245
MLAAttentionImpl,
246
)
@@ -234,7 +253,320 @@ from vllm.v1.attention.backends.utils import (
253
)
254
from vllm.v1.attention.ops.common import cp_lse_ag_out_rs
255
from vllm.v1.attention.ops.merge_attn_states import merge_attn_states
237
-from vllm.v1.kv_cache_interface import AttentionSpec
256
+from vllm.v1.attention.selector import get_attn_backend
257
+from vllm.v1.kv_cache_interface import (
258
+ AttentionSpec,
259
+ KVCacheSpec,
260
+ MLAAttentionSpec,
261
+)
262
+
263
+logger = init_logger(__name__)
264
+
265
+
266
+class MLAAttention(nn.Module, AttentionLayerBase):
267
+ """Multi-Head Latent Attention layer.
268
+
269
+ This class takes query, and compressed key/value tensors as input.
270
+ The class does the following:
271
+
272
+ 1. Store the input key and value tensors in the KV cache.
273
+ 2. Perform (multi-head/multi-query/grouped-query) attention.
274
+ 3. Return the output tensor.
275
+ """
276
+
277
+ def __init__(
278
+ self,
279
+ num_heads: int,
280
+ scale: float,
281
+ qk_nope_head_dim: int,
282
+ qk_rope_head_dim: int,
283
+ v_head_dim: int,
284
+ q_lora_rank: int | None,
285
+ kv_lora_rank: int,
286
+ kv_b_proj: ColumnParallelLinear,
287
+ cache_config: CacheConfig | None = None,
288
+ quant_config: QuantizationConfig | None = None,
289
+ prefix: str = "",
290
+ use_sparse: bool = False,
291
+ indexer: object | None = None,
292
+ **extra_impl_args,
293
+ ):
294
+ super().__init__()
295
+ self.num_heads = num_heads
296
+ self.scale = scale
297
+ self.qk_nope_head_dim = qk_nope_head_dim
298
+ self.qk_rope_head_dim = qk_rope_head_dim
299
+ self.v_head_dim = v_head_dim
300
+ self.q_lora_rank = q_lora_rank
301
+ self.kv_lora_rank = kv_lora_rank
302
+ self.head_size = kv_lora_rank + qk_rope_head_dim
303
+ self.layer_name = prefix
304
+
305
+ if cache_config is not None:
306
+ kv_cache_dtype = cache_config.cache_dtype
307
+ block_size = cache_config.block_size
308
+ calculate_kv_scales = cache_config.calculate_kv_scales
309
+ else:
310
+ kv_cache_dtype = "auto"
311
+ block_size = 16
312
+ calculate_kv_scales = False
313
+ self.quant_config = quant_config
314
+
315
+ # Initialize KV cache quantization attributes
316
+ self.kv_cache_dtype = kv_cache_dtype
317
+ self.calculate_kv_scales = calculate_kv_scales
318
+ _init_kv_cache_quant(self, quant_config, prefix)
319
+
320
+ dtype = torch.get_default_dtype()
321
+ self.attn_backend = get_attn_backend(
322
+ self.head_size,
323
+ dtype,
324
+ kv_cache_dtype,
325
+ block_size,
326
+ use_mla=True,
327
+ use_sparse=use_sparse,
328
+ )
329
+
330
+ if (
331
+ cache_config is not None
332
+ and cache_config.enable_prefix_caching
333
+ and vllm_is_batch_invariant()
334
+ and (
335
+ self.attn_backend.get_name() == "TRITON_MLA"
336
+ or self.attn_backend.get_name() == "FLASHINFER"
337
+ )
338
+ ):
339
+ logger.warning_once(
340
+ "Disabling prefix caching for TRITON_MLA / FLASHINFER "
341
+ "with batch invariance, as it is not yet supported.",
342
+ scope="local",
343
+ )
344
+ cache_config.enable_prefix_caching = False
345
+
346
+ impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
347
+ self.impl = impl_cls(
348
+ num_heads=self.num_heads,
349
+ head_size=self.head_size,
350
+ scale=self.scale,
351
+ num_kv_heads=1,
352
+ alibi_slopes=None,
353
+ sliding_window=None,
354
+ kv_cache_dtype=self.kv_cache_dtype,
355
+ logits_soft_cap=None,
356
+ attn_type=AttentionType.DECODER,
357
+ kv_sharing_target_layer_name=None,
358
+ # MLA Args
359
+ q_lora_rank=self.q_lora_rank,
360
+ kv_lora_rank=self.kv_lora_rank,
361
+ qk_nope_head_dim=self.qk_nope_head_dim,
362
+ qk_rope_head_dim=self.qk_rope_head_dim,
363
+ qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
364
+ v_head_dim=self.v_head_dim,
365
+ kv_b_proj=kv_b_proj,
366
+ indexer=indexer,
367
+ **extra_impl_args,
368
+ )
369
+
370
+ self.use_direct_call = not current_platform.opaque_attention_op()
371
+
372
+ compilation_config = get_current_vllm_config().compilation_config
373
+ if prefix in compilation_config.static_forward_context:
374
+ raise ValueError(f"Duplicate layer name: {prefix}")
375
+ compilation_config.static_forward_context[prefix] = self
376
+
377
+ self.kv_cache = [
378
+ torch.tensor([])
379
+ for _ in range(
380
+ get_current_vllm_config().parallel_config.pipeline_parallel_size
381
+ )
382
+ ]
383
+
384
+ self.use_sparse = use_sparse
385
+
386
+ # Initialize q/k/v range constants.
387
+ self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
388
+ self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
389
+ self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
390
+
391
+ def forward(
392
+ self,
393
+ q: torch.Tensor,
394
+ kv_c_normed: torch.Tensor,
395
+ k_pe: torch.Tensor,
396
+ output_shape: torch.Size | None = None,
397
+ ) -> torch.Tensor:
398
+ if self.calculate_kv_scales:
399
+ torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name)
400
+
401
+ if self.use_direct_call:
402
+ forward_context: ForwardContext = get_forward_context()
403
+ attn_metadata = forward_context.attn_metadata
404
+ if isinstance(attn_metadata, dict):
405
+ attn_metadata = attn_metadata[self.layer_name]
406
+ self_kv_cache = self.kv_cache[forward_context.virtual_engine]
407
+
408
+ if self.attn_backend.accept_output_buffer:
409
+ output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
410
+ self.impl.forward(
411
+ self,
412
+ q,
413
+ kv_c_normed,
414
+ k_pe,
415
+ self_kv_cache,
416
+ attn_metadata,
417
+ output=output,
418
+ )
419
+ return output
420
+ else:
421
+ return self.impl.forward(
422
+ self, q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
423
+ )
424
+ else:
425
+ if self.attn_backend.accept_output_buffer:
426
+ output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
427
+ torch.ops.vllm.unified_mla_attention_with_output(
428
+ q,
429
+ kv_c_normed,
430
+ k_pe,
431
+ output,
432
+ self.layer_name,
433
+ )
434
+ return output
435
+ else:
436
+ return torch.ops.vllm.unified_mla_attention(
437
+ q,
438
+ kv_c_normed,
439
+ k_pe,
440
+ self.layer_name,
441
+ )
442
+
443
+ def process_weights_after_loading(self, act_dtype: torch.dtype):
444
+ if hasattr(self.impl, "process_weights_after_loading"):
445
+ self.impl.process_weights_after_loading(act_dtype)
446
+
447
+ # If we should not load quant weights, we initialize the scales to 1.0
448
+ # as the default value. See [Note: Register q/k/v/prob scales in state dict]
449
+ # for more details.
450
+ quant_method = (
451
+ self.quant_config.get_quant_method(self, prefix=self.layer_name)
452
+ if self.quant_config
453
+ else None
454
+ )
455
+ if not should_load_quant_weights(quant_method):
456
+ set_default_quant_scales(self, register_buffer=False)
457
+
458
+ def calc_kv_scales(
459
+ self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
460
+ ) -> None:
461
+ """Optional scale calculation for MLA inputs.
462
+
463
+ Mirrors Attention.calc_kv_scales. Not all MLA backends require this
464
+ """
465
+ # Use safe defaults if ranges are not present
466
+ q_range = getattr(self, "q_range", torch.tensor(1.0))
467
+ k_range = getattr(self, "k_range", torch.tensor(1.0))
468
+ v_range = getattr(self, "v_range", torch.tensor(1.0))
469
+
470
+ self._q_scale.copy_(torch.abs(q).max() / q_range)
471
+ # kv_c_normed is the compressed KV representation; use it for k/v
472
+ kv_abs_max = torch.abs(kv_c_normed).max()
473
+ self._k_scale.copy_(kv_abs_max / k_range)
474
+ self._v_scale.copy_(kv_abs_max / v_range)
475
+ self._q_scale_float = self._q_scale.item()
476
+ self._k_scale_float = self._k_scale.item()
477
+ self._v_scale_float = self._v_scale.item()
478
+ self.calculate_kv_scales = False
479
+
480
+ def get_attn_backend(self) -> type[AttentionBackend]:
481
+ return self.attn_backend
482
+
483
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
484
+ kv_cache_dtype = kv_cache_dtype_str_to_dtype(
485
+ self.kv_cache_dtype, vllm_config.model_config
486
+ )
487
+ return MLAAttentionSpec(
488
+ block_size=vllm_config.cache_config.block_size,
489
+ num_kv_heads=1,
490
+ head_size=self.head_size,
491
+ dtype=kv_cache_dtype,
492
+ cache_dtype_str=vllm_config.cache_config.cache_dtype,
493
+ )
494
+
495
+
496
+@maybe_transfer_kv_layer
497
+def unified_mla_attention(
498
+ q: torch.Tensor,
499
+ kv_c_normed: torch.Tensor,
500
+ k_pe: torch.Tensor,
501
+ layer_name: str,
502
+) -> torch.Tensor:
503
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
504
+ output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)
505
+
506
+ return output
507
+
508
+
509
+def unified_mla_attention_fake(
510
+ q: torch.Tensor,
511
+ kv_c_normed: torch.Tensor,
512
+ k_pe: torch.Tensor,
513
+ layer_name: str,
514
+) -> torch.Tensor:
515
+ return torch.empty_like(q).contiguous()
516
+
517
+
518
+direct_register_custom_op(
519
+ op_name="unified_mla_attention",
520
+ op_func=unified_mla_attention,
521
+ mutates_args=[],
522
+ fake_impl=unified_mla_attention_fake,
523
+ dispatch_key=current_platform.dispatch_key,
524
+)
525
+
526
+
527
+@maybe_transfer_kv_layer
528
+def unified_mla_attention_with_output(
529
+ q: torch.Tensor,
530
+ kv_c_normed: torch.Tensor,
531
+ k_pe: torch.Tensor,
532
+ output: torch.Tensor,
533
+ layer_name: str,
534
+ output_scale: torch.Tensor | None = None,
535
+ output_block_scale: torch.Tensor | None = None,
536
+) -> None:
537
+ attn_metadata, self, kv_cache = get_attention_context(layer_name)
538
+ self.impl.forward(
539
+ self,
540
+ q,
541
+ kv_c_normed,
542
+ k_pe,
543
+ kv_cache,
544
+ attn_metadata,
545
+ output=output,
546
+ output_scale=output_scale,
547
+ output_block_scale=output_block_scale,
548
+ )
549
+
550
+
551
+def unified_mla_attention_with_output_fake(
552
+ q: torch.Tensor,
553
+ kv_c_normed: torch.Tensor,
554
+ k_pe: torch.Tensor,
555
+ output: torch.Tensor,
556
+ layer_name: str,
557
+ output_scale: torch.Tensor | None = None,
558
+ output_block_scale: torch.Tensor | None = None,
559
+) -> None:
560
+ return
561
+
562
+
563
+direct_register_custom_op(
564
+ op_name="unified_mla_attention_with_output",
565
+ op_func=unified_mla_attention_with_output,
566
+ mutates_args=["output", "output_block_scale"],
567
+ fake_impl=unified_mla_attention_with_output_fake,
568
+ dispatch_key=current_platform.dispatch_key,
569
+)
570
571
572
class QueryLenSupport(Enum):
@@ -266,15 +598,12 @@ except ImportError:
598
from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]
599
is_vllm_fa = False
600
269
-try:
270
- from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
271
- from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache # noqa: F401
601
273
- flashinfer_available = True
274
-except ImportError:
275
- BatchPrefillWithRaggedKVCacheWrapper = object
602
+@functools.cache
603
+def flashinfer_available() -> bool:
604
+ import importlib.util
605
277
- flashinfer_available = False
606
+ return importlib.util.find_spec("flashinfer") is not None
607
608
609
def dynamic_per_batched_tensor_quant(
@@ -398,8 +727,8 @@ class MLACommonPrefillMetadata:
727
728
@dataclass
729
class FlashInferPrefillMetadata(MLACommonPrefillMetadata):
401
- prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None
402
- prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = field(
730
+ prefill_main: "BatchPrefillWithRaggedKVCacheWrapper | None" = None
731
+ prefill_chunks: "list[BatchPrefillWithRaggedKVCacheWrapper]" = field(
732
default_factory=list
733
)
734
@@ -495,7 +824,7 @@ def use_flashinfer_prefill() -> bool:
824
vllm_config = get_current_vllm_config()
825
if not (
826
not vllm_config.attention_config.disable_flashinfer_prefill
498
- and flashinfer_available
827
+ and flashinfer_available()
828
and not vllm_config.attention_config.use_cudnn_prefill
829
and current_platform.is_device_capability_family(100)
830
):
@@ -509,7 +838,7 @@ def use_cudnn_prefill() -> bool:
838
839
vllm_config = get_current_vllm_config()
840
return (
512
- flashinfer_available
841
+ flashinfer_available()
842
and vllm_config.attention_config.use_cudnn_prefill
843
and current_platform.is_device_capability_family(100)
844
and has_nvidia_artifactory()
@@ -731,6 +1060,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1060
has_context = True
1061
1062
if self._fi_prefill_main is None:
1063
+ from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
1064
+
1065
self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper(
1066
self._workspace_buffer, "NHD", backend="cutlass"
1067
)
@@ -739,6 +1070,8 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1070
num_chunks = chunked_context.cu_seq_lens.shape[0]
1071
# Allocate more prefill chunk wrappers if needed
1072
if len(self._fi_prefill_chunks) < num_chunks:
1073
+ from flashinfer import BatchPrefillWithRaggedKVCacheWrapper
1074
+
1075
for _ in range(len(self._fi_prefill_chunks), num_chunks):
1076
self._fi_prefill_chunks.append(
1077
BatchPrefillWithRaggedKVCacheWrapper(
@@ -1513,6 +1846,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1846
):
1847
assert isinstance(prefill, CudnnPrefillMetadata)
1848
assert prefill.query_seq_lens is not None
1849
+ from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
1850
+
1851
output, lse = cudnn_batch_prefill_with_kv_cache(
1852
q=q,
1853
k_cache=k,
@@ -1572,6 +1907,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1907
assert prefill.chunked_context is not None
1908
assert prefill.chunked_context.seq_lens[chunk_idx] is not None
1909
assert prefill.query_seq_lens is not None
1910
+ from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache
1911
+
1912
return cudnn_batch_prefill_with_kv_cache(
1913
q=q,
1914
k_cache=k,
vllm/model_executor/layers/attention/static_sink_attention.py
+1
-1
@@ -4,11 +4,11 @@ import functools
4
5
import torch
6
7
-from vllm.attention.layer import Attention
7
from vllm.config import CacheConfig, VllmConfig
8
from vllm.forward_context import ForwardContext, get_forward_context
9
from vllm.logger import init_logger
10
from vllm.model_executor.custom_op import CustomOp
11
+from vllm.model_executor.layers.attention import Attention
12
from vllm.utils.math_utils import cdiv
13
from vllm.utils.torch_utils import direct_register_custom_op
14
from vllm.v1.attention.backend import (
vllm/model_executor/layers/mla.py
+1
-1
@@ -4,9 +4,9 @@ from dataclasses import dataclass
4
5
import torch
6
7
-from vllm.attention.layer import MLAAttention
7
from vllm.config import CacheConfig
8
from vllm.model_executor.custom_op import PluggableLayer
9
+from vllm.model_executor.layers.attention import MLAAttention
10
from vllm.model_executor.layers.quantization import QuantizationConfig
11
12
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+1
-1
@@ -19,12 +19,12 @@ from compressed_tensors.quantization import (
19
from compressed_tensors.transform import TransformConfig
20
21
import vllm.envs as envs
22
-from vllm.attention.layer import Attention
22
from vllm.distributed import (
23
get_tensor_model_parallel_rank,
24
get_tensor_model_parallel_world_size,
25
)
26
from vllm.logger import init_logger
27
+from vllm.model_executor.layers.attention import Attention
28
from vllm.model_executor.layers.fused_moe import FusedMoE
29
from vllm.model_executor.layers.linear import (
30
LinearBase,
vllm/model_executor/layers/quantization/fp8.py
+1
-1
@@ -11,9 +11,9 @@ import vllm.envs as envs
11
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
12
from vllm import _custom_ops as ops
13
from vllm._aiter_ops import rocm_aiter_ops
14
-from vllm.attention.layer import Attention
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.batch_invariant import (
18
vllm_is_batch_invariant,
19
)
vllm/model_executor/layers/quantization/modelopt.py
+1
-1
@@ -11,8 +11,8 @@ from torch.nn.parameter import Parameter
11
import vllm.envs as envs
12
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
13
from vllm._custom_ops import cutlass_scaled_fp4_mm, scaled_fp4_quant
14
-from vllm.attention.layer import Attention
14
from vllm.logger import init_logger
15
+from vllm.model_executor.layers.attention import Attention
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEConfig,
18
FusedMoEQuantConfig,
vllm/model_executor/layers/quantization/mxfp4.py
+1
-1
@@ -7,9 +7,9 @@ import torch
7
from torch.nn.parameter import Parameter
8
9
from vllm import envs
10
-from vllm.attention.layer import Attention
10
from vllm.config import get_current_vllm_config
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention import Attention
13
from vllm.model_executor.layers.fused_moe import (
14
FusedMoE,
15
FusedMoEConfig,
vllm/model_executor/layers/quantization/petit.py
+1
-1
@@ -8,8 +8,8 @@ import regex as re
8
import torch
9
from torch.nn.parameter import Parameter
10
11
-from vllm.attention.layer import Attention
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention import Attention
13
from vllm.model_executor.layers.linear import (
14
LinearBase,
15
LinearMethodBase,
vllm/model_executor/layers/quantization/ptpc_fp8.py
+1
-1
@@ -7,8 +7,8 @@ import torch
7
from torch.nn.parameter import Parameter
8
9
from vllm import _custom_ops as ops
10
-from vllm.attention.layer import Attention
10
from vllm.logger import init_logger
11
+from vllm.model_executor.layers.attention import Attention
12
from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
13
from vllm.model_executor.layers.quantization import QuantizationMethods
14
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
vllm/model_executor/layers/quantization/quark/quark.py
+1
-1
@@ -6,8 +6,8 @@ from typing import TYPE_CHECKING, Any, Optional, cast
6
7
import torch
8
9
-from vllm.attention.layer import Attention
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.attention import Attention
11
from vllm.model_executor.layers.fused_moe import FusedMoE
12
from vllm.model_executor.layers.linear import (
13
LinearBase,
vllm/model_executor/model_loader/utils.py
+1
-1
@@ -11,9 +11,9 @@ import torch
11
from torch import nn
12
from typing_extensions import assert_never
13
14
-from vllm.attention.layer import Attention, MLAAttention
14
from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.attention import Attention, MLAAttention
17
from vllm.model_executor.layers.quantization.base_config import (
18
QuantizationConfig,
19
QuantizeMethodBase,
vllm/model_executor/models/afmoe.py
+1
-1
@@ -9,7 +9,6 @@ from itertools import islice
9
import torch
10
from torch import nn
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
14
from vllm.distributed import (
@@ -18,6 +17,7 @@ from vllm.distributed import (
17
get_tensor_model_parallel_world_size,
18
)
19
from vllm.logger import init_logger
20
+from vllm.model_executor.layers.attention import Attention
21
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
22
from vllm.model_executor.layers.layernorm import RMSNorm
23
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/aimv2.py
+1
-1
@@ -11,7 +11,7 @@ import torch.nn as nn
11
from vllm.distributed import get_tensor_model_parallel_world_size
12
from vllm.distributed.utils import divide
13
from vllm.model_executor.layers.activation import SiluAndMul
14
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
14
+from vllm.model_executor.layers.attention import MMEncoderAttention
15
from vllm.model_executor.layers.conv import Conv2dLayer
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/apertus.py
+2
-2
@@ -32,12 +32,12 @@ import torch
32
from torch import nn
33
from transformers import ApertusConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.model_executor.layers.activation import XIELU
40
-from vllm.model_executor.layers.attention.encoder_only_attention import (
39
+from vllm.model_executor.layers.attention import (
40
+ Attention,
41
EncoderOnlyAttention,
42
)
43
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/arctic.py
+1
-1
@@ -8,7 +8,6 @@ from itertools import islice
8
import torch
9
from torch import nn
10
11
-from vllm.attention.layer import Attention
11
from vllm.compilation.decorators import support_torch_compile
12
from vllm.config import CacheConfig, VllmConfig
13
from vllm.distributed import (
@@ -19,6 +18,7 @@ from vllm.distributed import (
18
)
19
from vllm.logger import init_logger
20
from vllm.model_executor.layers.activation import SiluAndMul
21
+from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/baichuan.py
+1
-1
@@ -29,7 +29,6 @@ import torch
29
from torch import nn
30
from transformers import PretrainedConfig
31
32
-from vllm.attention.layer import Attention
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import (
@@ -38,6 +37,7 @@ from vllm.distributed import (
37
get_tensor_model_parallel_world_size,
38
)
39
from vllm.model_executor.layers.activation import SiluAndMul
40
+from vllm.model_executor.layers.attention import Attention
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import (
43
MergedColumnParallelLinear,
vllm/model_executor/models/bailing_moe.py
+1
-1
@@ -32,7 +32,6 @@ import torch.nn.functional as F
32
from torch import nn
33
from transformers.configuration_utils import PretrainedConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import (
@@ -41,6 +40,7 @@ from vllm.distributed import (
40
get_tensor_model_parallel_world_size,
41
)
42
from vllm.model_executor.layers.activation import SiluAndMul
43
+from vllm.model_executor.layers.attention import Attention
44
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/bamba.py
+1
-1
@@ -9,12 +9,12 @@ import torch
9
from torch import nn
10
from transformers import BambaConfig
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.model_executor.layers.activation import SiluAndMul
17
+from vllm.model_executor.layers.attention import Attention
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (
20
MergedColumnParallelLinear,
vllm/model_executor/models/bert.py
+1
-1
@@ -11,7 +11,7 @@ from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import CacheConfig, ModelConfig, PoolerConfig, VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import get_act_fn
14
-from vllm.model_executor.layers.attention.encoder_only_attention import (
14
+from vllm.model_executor.layers.attention import (
15
EncoderOnlyAttention,
16
)
17
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/bert_with_rope.py
+1
-1
@@ -15,7 +15,7 @@ from vllm.distributed import (
15
tensor_model_parallel_all_reduce,
16
)
17
from vllm.model_executor.layers.activation import get_act_and_mul_fn, get_act_fn
18
-from vllm.model_executor.layers.attention.encoder_only_attention import (
18
+from vllm.model_executor.layers.attention import (
19
EncoderOnlyAttention,
20
)
21
from vllm.model_executor.layers.fused_moe import activation_without_mul, fused_topk
vllm/model_executor/models/blip.py
+1
-1
@@ -11,7 +11,7 @@ from transformers import Blip2VisionConfig, BlipVisionConfig
11
12
from vllm.distributed import divide, get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import get_act_fn
14
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
14
+from vllm.model_executor.layers.attention import MMEncoderAttention
15
from vllm.model_executor.layers.conv import Conv2dLayer
16
from vllm.model_executor.layers.linear import (
17
ColumnParallelLinear,
vllm/model_executor/models/bloom.py
+1
-1
@@ -27,7 +27,6 @@ import torch
27
from torch import nn
28
from transformers import BloomConfig
29
30
-from vllm.attention.layer import Attention
30
from vllm.compilation.decorators import support_torch_compile
31
from vllm.config import CacheConfig, VllmConfig
32
from vllm.distributed import (
@@ -36,6 +35,7 @@ from vllm.distributed import (
35
get_tensor_model_parallel_world_size,
36
)
37
from vllm.model_executor.layers.activation import get_act_fn
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.linear import (
40
ColumnParallelLinear,
41
QKVParallelLinear,
vllm/model_executor/models/chameleon.py
+1
-1
@@ -16,12 +16,12 @@ from transformers import (
16
ChameleonVQVAEConfig,
17
)
18
19
-from vllm.attention.layer import Attention
19
from vllm.config import CacheConfig, VllmConfig
20
from vllm.config.multimodal import BaseDummyOptions
21
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.activation import SiluAndMul
24
+from vllm.model_executor.layers.attention import Attention
25
from vllm.model_executor.layers.conv import Conv2dLayer
26
from vllm.model_executor.layers.layernorm import RMSNorm
27
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/chatglm.py
+1
-1
@@ -12,11 +12,11 @@ import torch
12
from torch import nn
13
from torch.nn import LayerNorm
14
15
-from vllm.attention.layer import Attention
15
from vllm.compilation.decorators import support_torch_compile
16
from vllm.config import CacheConfig, VllmConfig
17
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
18
from vllm.model_executor.layers.activation import SiluAndMul
19
+from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import (
22
MergedColumnParallelLinear,
vllm/model_executor/models/clip.py
+1
-2
@@ -14,12 +14,11 @@ from transformers import (
14
CLIPVisionConfig,
15
)
16
17
-from vllm.attention.layer import Attention
17
from vllm.config import VllmConfig
18
from vllm.config.multimodal import BaseDummyOptions
19
from vllm.distributed import divide, get_tensor_model_parallel_world_size
20
from vllm.model_executor.layers.activation import get_act_fn
22
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
21
+from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
22
from vllm.model_executor.layers.conv import Conv2dLayer
23
from vllm.model_executor.layers.linear import (
24
ColumnParallelLinear,
vllm/model_executor/models/commandr.py
+1
-1
@@ -30,11 +30,11 @@ import torch
30
from torch import nn
31
from transformers import Cohere2Config, CohereConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import SiluAndMul
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.linear import (
39
MergedColumnParallelLinear,
40
QKVParallelLinear,
vllm/model_executor/models/dbrx.py
+1
-1
@@ -8,13 +8,13 @@ import torch
8
import torch.nn as nn
9
from transformers import DbrxConfig
10
11
-from vllm.attention.layer import Attention
11
from vllm.config import CacheConfig, VllmConfig
12
from vllm.distributed import (
13
get_pp_group,
14
get_tensor_model_parallel_rank,
15
get_tensor_model_parallel_world_size,
16
)
17
+from vllm.model_executor.layers.attention import Attention
18
from vllm.model_executor.layers.fused_moe import FusedMoE
19
from vllm.model_executor.layers.linear import (
20
QKVParallelLinear,
vllm/model_executor/models/deepencoder.py
+1
-1
@@ -18,7 +18,7 @@ import torch.nn as nn
18
import torch.nn.functional as F
19
from transformers import CLIPVisionConfig
20
21
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
21
+from vllm.model_executor.layers.attention import MMEncoderAttention
22
from vllm.model_executor.layers.conv import Conv2dLayer
23
from vllm.model_executor.layers.quantization import QuantizationConfig
24
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
vllm/model_executor/models/deepseek_v2.py
+1
-1
@@ -33,7 +33,6 @@ from torch import nn
33
from transformers import DeepseekV2Config, DeepseekV3Config
34
35
from vllm._aiter_ops import rocm_aiter_ops
36
-from vllm.attention.layer import Attention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, ParallelConfig, VllmConfig, get_current_vllm_config
38
from vllm.distributed import (
@@ -45,6 +44,7 @@ from vllm.distributed import (
44
)
45
from vllm.logger import init_logger
46
from vllm.model_executor.layers.activation import SiluAndMul
47
+from vllm.model_executor.layers.attention import Attention
48
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
49
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
50
from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
vllm/model_executor/models/dots1.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
from torch import nn
33
from transformers import Dots1Config
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, ModelConfig, VllmConfig
37
from vllm.distributed import (
@@ -41,6 +40,7 @@ from vllm.distributed import (
40
tensor_model_parallel_all_reduce,
41
)
42
from vllm.model_executor.layers.activation import SiluAndMul
43
+from vllm.model_executor.layers.attention import Attention
44
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
45
from vllm.model_executor.layers.layernorm import RMSNorm
46
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/dots_ocr.py
+1
-1
@@ -16,7 +16,7 @@ from vllm.distributed.parallel_state import (
16
get_tensor_model_parallel_world_size,
17
)
18
from vllm.model_executor.layers.activation import SiluAndMul
19
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
19
+from vllm.model_executor.layers.attention import (
20
MMEncoderAttention,
21
)
22
from vllm.model_executor.layers.conv import Conv2dLayer
vllm/model_executor/models/ernie45_moe.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
from torch import nn
33
from transformers import PretrainedConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
37
from vllm.distributed import (
@@ -42,6 +41,7 @@ from vllm.distributed import (
41
)
42
from vllm.logger import init_logger
43
from vllm.model_executor.layers.activation import SiluAndMul
44
+from vllm.model_executor.layers.attention import Attention
45
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/ernie45_vl.py
+1
-1
@@ -42,7 +42,7 @@ from vllm.distributed import parallel_state
42
from vllm.distributed import utils as dist_utils
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import QuickGELU
45
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
45
+from vllm.model_executor.layers.attention import (
46
MMEncoderAttention,
47
)
48
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/ernie45_vl_moe.py
+1
-2
@@ -31,12 +31,11 @@ import torch
31
from torch import nn
32
from transformers import PretrainedConfig
33
34
-from vllm.attention.layer import Attention
35
-
34
# from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.logger import init_logger
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
40
from vllm.model_executor.layers.layernorm import RMSNorm
41
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/exaone.py
+1
-1
@@ -32,11 +32,11 @@ import torch
32
from torch import nn
33
from transformers import PretrainedConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.model_executor.layers.activation import SiluAndMul
39
+from vllm.model_executor.layers.attention import Attention
40
from vllm.model_executor.layers.layernorm import RMSNorm
41
from vllm.model_executor.layers.linear import (
42
MergedColumnParallelLinear,
vllm/model_executor/models/exaone4.py
+1
-1
@@ -28,11 +28,11 @@ import torch
28
from torch import nn
29
from transformers import Exaone4Config
30
31
-from vllm.attention.layer import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.model_executor.layers.activation import SiluAndMul
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.layernorm import RMSNorm
37
from vllm.model_executor.layers.linear import (
38
MergedColumnParallelLinear,
vllm/model_executor/models/falcon.py
+1
-1
@@ -30,7 +30,6 @@ from torch import nn
30
from torch.nn import LayerNorm
31
from transformers import FalconConfig as HF_FalconConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import (
@@ -40,6 +39,7 @@ from vllm.distributed import (
39
tensor_model_parallel_all_reduce,
40
)
41
from vllm.model_executor.layers.activation import get_act_fn
42
+from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.linear import (
44
ColumnParallelLinear,
45
QKVParallelLinear,
vllm/model_executor/models/falcon_h1.py
+1
-1
@@ -9,12 +9,12 @@ import torch
9
from torch import nn
10
from transformers import FalconH1Config
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
from vllm.model_executor.layers.activation import SiluAndMul
17
+from vllm.model_executor.layers.attention import Attention
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (
20
MergedColumnParallelLinear,
vllm/model_executor/models/gemma.py
+1
-1
@@ -26,12 +26,12 @@ import torch
26
from torch import nn
27
from transformers import GemmaConfig
28
29
-from vllm.attention.layer import Attention
29
from vllm.compilation.decorators import support_torch_compile
30
from vllm.config import CacheConfig, VllmConfig
31
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
32
from vllm.logger import init_logger
33
from vllm.model_executor.layers.activation import GeluAndMul
34
+from vllm.model_executor.layers.attention import Attention
35
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
36
from vllm.model_executor.layers.linear import (
37
MergedColumnParallelLinear,
vllm/model_executor/models/gemma2.py
+1
-1
@@ -23,12 +23,12 @@ import torch
23
from torch import nn
24
from transformers import Gemma2Config
25
26
-from vllm.attention.layer import Attention
26
from vllm.compilation.decorators import support_torch_compile
27
from vllm.config import CacheConfig, VllmConfig
28
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
29
from vllm.logger import init_logger
30
from vllm.model_executor.layers.activation import GeluAndMul
31
+from vllm.model_executor.layers.attention import Attention
32
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
33
from vllm.model_executor.layers.linear import (
34
MergedColumnParallelLinear,
vllm/model_executor/models/gemma3.py
+2
-2
@@ -22,13 +22,13 @@ import torch
22
from torch import nn
23
from transformers import Gemma3TextConfig
24
25
-from vllm.attention.layer import Attention
25
from vllm.compilation.decorators import support_torch_compile
26
from vllm.config import CacheConfig, VllmConfig
27
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
28
from vllm.logger import init_logger
29
from vllm.model_executor.layers.activation import GeluAndMul
31
-from vllm.model_executor.layers.attention.encoder_only_attention import (
30
+from vllm.model_executor.layers.attention import (
31
+ Attention,
32
EncoderOnlyAttention,
33
)
34
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
vllm/model_executor/models/gemma3n.py
+1
-1
@@ -21,7 +21,6 @@ import torch
21
from torch import nn
22
from transformers.models.gemma3n.configuration_gemma3n import Gemma3nTextConfig
23
24
-from vllm.attention.layer import Attention
24
from vllm.compilation.decorators import support_torch_compile
25
from vllm.config import CacheConfig, VllmConfig
26
from vllm.distributed import get_tensor_model_parallel_world_size
@@ -32,6 +31,7 @@ from vllm.model_executor.layers.activation import (
31
GeluAndMul,
32
GeluAndMulSparse,
33
)
34
+from vllm.model_executor.layers.attention import Attention
35
from vllm.model_executor.layers.layernorm import RMSNorm
36
from vllm.model_executor.layers.linear import (
37
ColumnParallelLinear,
vllm/model_executor/models/glm4.py
+1
-1
@@ -29,10 +29,10 @@ import torch
29
from torch import nn
30
from transformers import Glm4Config
31
32
-from vllm.attention.layer import Attention
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, VllmConfig
34
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.layernorm import RMSNorm
37
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
38
from vllm.model_executor.layers.logits_processor import LogitsProcessor
vllm/model_executor/models/glm4_1v.py
+1
-1
@@ -52,7 +52,7 @@ from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
52
from vllm.distributed import get_tensor_model_parallel_world_size, parallel_state
53
from vllm.distributed import utils as dist_utils
54
from vllm.logger import init_logger
55
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
55
+from vllm.model_executor.layers.attention import (
56
MMEncoderAttention,
57
)
58
from vllm.model_executor.layers.conv import Conv2dLayer, Conv3dLayer
vllm/model_executor/models/glm4_moe.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
from torch import nn
33
from transformers.models.glm4_moe import Glm4MoeConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
37
from vllm.distributed import (
@@ -42,6 +41,7 @@ from vllm.distributed import (
41
)
42
from vllm.logger import init_logger
43
from vllm.model_executor.layers.activation import SiluAndMul
44
+from vllm.model_executor.layers.attention import Attention
45
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
46
from vllm.model_executor.layers.layernorm import RMSNorm
47
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/glm4v.py
+1
-1
@@ -23,7 +23,7 @@ from vllm.config import VllmConfig
23
from vllm.config.multimodal import BaseDummyOptions
24
from vllm.distributed import get_tensor_model_parallel_world_size
25
from vllm.model_executor.layers.activation import SiluAndMul, get_act_fn
26
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
26
+from vllm.model_executor.layers.attention import MMEncoderAttention
27
from vllm.model_executor.layers.conv import Conv2dLayer
28
from vllm.model_executor.layers.linear import (
29
ColumnParallelLinear,
vllm/model_executor/models/glmasr.py
+1
-1
@@ -16,7 +16,7 @@ from vllm.config.multimodal import BaseDummyOptions
16
from vllm.distributed.parallel_state import get_tensor_model_parallel_world_size
17
from vllm.inputs.data import PromptType
18
from vllm.model_executor.layers.activation import get_act_fn
19
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
19
+from vllm.model_executor.layers.attention import MMEncoderAttention
20
from vllm.model_executor.layers.linear import (
21
ColumnParallelLinear,
22
QKVParallelLinear,
vllm/model_executor/models/gpt2.py
+1
-1
@@ -27,7 +27,6 @@ import torch
27
from torch import nn
28
from transformers import GPT2Config
29
30
-from vllm.attention.layer import Attention
30
from vllm.compilation.decorators import support_torch_compile
31
from vllm.config import CacheConfig, VllmConfig
32
from vllm.distributed.parallel_state import (
@@ -35,6 +34,7 @@ from vllm.distributed.parallel_state import (
34
get_tensor_model_parallel_world_size,
35
)
36
from vllm.model_executor.layers.activation import get_act_fn
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.linear import (
39
ColumnParallelLinear,
40
QKVParallelLinear,
vllm/model_executor/models/gpt_bigcode.py
+1
-1
@@ -28,11 +28,11 @@ import torch
28
from torch import nn
29
from transformers import GPTBigCodeConfig
30
31
-from vllm.attention.layer import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.model_executor.layers.activation import get_act_fn
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.linear import (
37
ColumnParallelLinear,
38
QKVParallelLinear,
vllm/model_executor/models/gpt_j.py
+1
-1
@@ -26,11 +26,11 @@ import torch
26
from torch import nn
27
from transformers import GPTJConfig
28
29
-from vllm.attention.layer import Attention
29
from vllm.compilation.decorators import support_torch_compile
30
from vllm.config import CacheConfig, VllmConfig
31
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
32
from vllm.model_executor.layers.activation import get_act_fn
33
+from vllm.model_executor.layers.attention import Attention
34
from vllm.model_executor.layers.linear import (
35
ColumnParallelLinear,
36
QKVParallelLinear,
vllm/model_executor/models/gpt_neox.py
+1
-1
@@ -26,11 +26,11 @@ import torch
26
from torch import nn
27
from transformers import GPTNeoXConfig
28
29
-from vllm.attention.layer import Attention
29
from vllm.compilation.decorators import support_torch_compile
30
from vllm.config import CacheConfig, VllmConfig
31
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
32
from vllm.model_executor.layers.activation import get_act_fn
33
+from vllm.model_executor.layers.attention import Attention
34
from vllm.model_executor.layers.linear import (
35
ColumnParallelLinear,
36
QKVParallelLinear,
vllm/model_executor/models/gpt_oss.py
+1
-1
@@ -7,7 +7,6 @@ import torch.distributed as dist
7
from torch import nn
8
from transformers import GptOssConfig
9
10
-from vllm.attention.layer import Attention
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import CacheConfig, VllmConfig
12
from vllm.distributed import (
@@ -19,6 +18,7 @@ from vllm.distributed import (
18
get_tensor_model_parallel_world_size,
19
tensor_model_parallel_all_gather,
20
)
21
+from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.fused_moe import FusedMoE
23
from vllm.model_executor.layers.fused_moe.config import FusedMoEParallelConfig
24
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/granite.py
+1
-1
@@ -31,11 +31,11 @@ import torch
31
from torch import nn
32
from transformers import GraniteConfig
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
MergedColumnParallelLinear,
vllm/model_executor/models/granitemoe.py
+1
-1
@@ -31,7 +31,6 @@ from typing import Any
31
import torch
32
from torch import nn
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import (
@@ -39,6 +38,7 @@ from vllm.distributed import (
38
get_tensor_model_parallel_world_size,
39
tensor_model_parallel_all_gather,
40
)
41
+from vllm.model_executor.layers.attention import Attention
42
from vllm.model_executor.layers.fused_moe import FusedMoE
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/granitemoehybrid.py
+1
-1
@@ -9,11 +9,11 @@ import torch
9
from torch import nn
10
from transformers import GraniteMoeHybridConfig
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
+from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.layernorm import RMSNorm
18
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
19
from vllm.model_executor.layers.logits_processor import LogitsProcessor
vllm/model_executor/models/grok1.py
+1
-1
@@ -32,12 +32,12 @@ import torch
32
import torch.nn.functional as F
33
from torch import nn
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
38
from vllm.logger import init_logger
39
from vllm.model_executor.layers.activation import GeluAndMul
40
+from vllm.model_executor.layers.attention import Attention
41
from vllm.model_executor.layers.fused_moe import FusedMoE
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/hunyuan_v1.py
+1
-1
@@ -33,7 +33,6 @@ import torch
33
from torch import nn
34
from transformers import PretrainedConfig
35
36
-from vllm.attention.layer import Attention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
38
from vllm.distributed import (
@@ -43,6 +42,7 @@ from vllm.distributed import (
42
tensor_model_parallel_all_reduce,
43
)
44
from vllm.model_executor.layers.activation import SiluAndMul
45
+from vllm.model_executor.layers.attention import Attention
46
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/hunyuan_vision.py
+1
-1
@@ -39,7 +39,7 @@ from vllm.distributed import parallel_state
39
from vllm.distributed import utils as dist_utils
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.activation import get_act_fn
42
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
42
+from vllm.model_executor.layers.attention import MMEncoderAttention
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
ColumnParallelLinear,
vllm/model_executor/models/idefics2_vision_model.py
+1
-1
@@ -29,7 +29,7 @@ from transformers.models.idefics2.configuration_idefics2 import (
29
30
from vllm.distributed import get_tensor_model_parallel_world_size
31
from vllm.model_executor.layers.activation import get_act_fn
32
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
32
+from vllm.model_executor.layers.attention import MMEncoderAttention
33
from vllm.model_executor.layers.conv import Conv2dLayer
34
from vllm.model_executor.layers.linear import (
35
ColumnParallelLinear,
vllm/model_executor/models/intern_vit.py
+1
-1
@@ -23,7 +23,7 @@ from vllm.distributed import (
23
tensor_model_parallel_all_gather,
24
)
25
from vllm.model_executor.layers.activation import get_act_fn
26
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
26
+from vllm.model_executor.layers.attention import MMEncoderAttention
27
from vllm.model_executor.layers.conv import Conv2dLayer
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/internlm2.py
+1
-1
@@ -10,7 +10,6 @@ import torch
10
from torch import nn
11
from transformers import PretrainedConfig
12
13
-from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, VllmConfig
15
from vllm.distributed import (
@@ -21,6 +20,7 @@ from vllm.distributed import (
20
tensor_model_parallel_all_gather,
21
)
22
from vllm.model_executor.layers.activation import SiluAndMul
23
+from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
26
MergedColumnParallelLinear,
vllm/model_executor/models/interns1_vit.py
+1
-1
@@ -15,7 +15,7 @@ from transformers import PretrainedConfig
15
from transformers.utils import torch_int
16
17
from vllm.model_executor.layers.activation import get_act_fn
18
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
18
+from vllm.model_executor.layers.attention import MMEncoderAttention
19
from vllm.model_executor.layers.conv import Conv2dLayer
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
vllm/model_executor/models/iquest_loopcoder.py
+1
-1
@@ -24,10 +24,10 @@ import torch
24
from torch import nn
25
from transformers import PretrainedConfig
26
27
-from vllm.attention.layer import Attention
27
from vllm.compilation.decorators import support_torch_compile
28
from vllm.config import CacheConfig, VllmConfig
29
from vllm.distributed import get_tensor_model_parallel_world_size
30
+from vllm.model_executor.layers.attention import Attention
31
from vllm.model_executor.layers.layernorm import RMSNorm
32
from vllm.model_executor.layers.linear import (
33
ColumnParallelLinear,
vllm/model_executor/models/isaac.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.config import VllmConfig
20
from vllm.config.model import ModelConfig
21
from vllm.distributed import parallel_state
22
from vllm.distributed import utils as dist_utils
23
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
23
+from vllm.model_executor.layers.attention import MMEncoderAttention
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
26
QKVParallelLinear,
vllm/model_executor/models/jais.py
+1
-1
@@ -28,7 +28,6 @@ from itertools import islice
28
import torch
29
from torch import nn
30
31
-from vllm.attention.layer import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import (
@@ -36,6 +35,7 @@ from vllm.distributed import (
35
get_tensor_model_parallel_rank,
36
get_tensor_model_parallel_world_size,
37
)
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.linear import (
40
ColumnParallelLinear,
41
QKVParallelLinear,
vllm/model_executor/models/jais2.py
+1
-1
@@ -31,7 +31,6 @@ import torch
31
from torch import nn
32
from transformers import Jais2Config
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import (
@@ -39,6 +38,7 @@ from vllm.distributed import (
38
get_tensor_model_parallel_world_size,
39
)
40
from vllm.model_executor.layers.activation import ReLUSquaredActivation
41
+from vllm.model_executor.layers.attention import Attention
42
from vllm.model_executor.layers.linear import (
43
ColumnParallelLinear,
44
QKVParallelLinear,
vllm/model_executor/models/jamba.py
+1
-1
@@ -9,11 +9,11 @@ import torch
9
from torch import nn
10
from transformers import JambaConfig
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.distributed.parallel_state import get_pp_group
16
+from vllm.model_executor.layers.attention import Attention
17
from vllm.model_executor.layers.fused_moe import FusedMoE
18
from vllm.model_executor.layers.layernorm import RMSNorm
19
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/keye.py
+1
-1
@@ -20,7 +20,7 @@ from vllm.config import VllmConfig
20
from vllm.config.multimodal import BaseDummyOptions
21
from vllm.distributed import get_tensor_model_parallel_world_size
22
from vllm.logger import init_logger
23
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
23
+from vllm.model_executor.layers.attention import (
24
MMEncoderAttention,
25
)
26
from vllm.model_executor.layers.conv import Conv2dLayer
vllm/model_executor/models/lfm2.py
+1
-1
@@ -7,11 +7,11 @@ import torch
7
import torch.nn as nn
8
from transformers import Lfm2Config
9
10
-from vllm.attention.layer import Attention
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import CacheConfig, ModelConfig, VllmConfig
12
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
13
from vllm.model_executor.layers.activation import SiluAndMul
14
+from vllm.model_executor.layers.attention import Attention
15
from vllm.model_executor.layers.layernorm import RMSNorm
16
from vllm.model_executor.layers.linear import (
17
MergedColumnParallelLinear,
vllm/model_executor/models/lfm2_moe.py
+1
-1
@@ -6,7 +6,6 @@ from itertools import islice
6
import torch
7
import torch.nn as nn
8
9
-from vllm.attention.layer import Attention
9
from vllm.compilation.decorators import support_torch_compile
10
from vllm.config import CacheConfig, ModelConfig, VllmConfig, get_current_vllm_config
11
from vllm.distributed import (
@@ -15,6 +14,7 @@ from vllm.distributed import (
14
get_tensor_model_parallel_world_size,
15
)
16
from vllm.model_executor.layers.activation import SiluAndMul
17
+from vllm.model_executor.layers.attention import Attention
18
from vllm.model_executor.layers.fused_moe import FusedMoE
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/lfm2_siglip2.py
+1
-1
@@ -13,7 +13,7 @@ from transformers import Siglip2VisionConfig
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
from vllm.model_executor.layers.activation import get_act_fn
16
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
16
+from vllm.model_executor.layers.attention import MMEncoderAttention
17
from vllm.model_executor.layers.linear import (
18
ColumnParallelLinear,
19
QKVParallelLinear,
vllm/model_executor/models/llama.py
+2
-2
@@ -31,12 +31,12 @@ import torch
31
from torch import nn
32
from transformers import LlamaConfig
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.model_executor.layers.activation import SiluAndMul
39
-from vllm.model_executor.layers.attention.encoder_only_attention import (
38
+from vllm.model_executor.layers.attention import (
39
+ Attention,
40
EncoderOnlyAttention,
41
)
42
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/llama4.py
+2
-2
@@ -24,7 +24,6 @@ import torch
24
from torch import nn
25
from transformers import Llama4TextConfig
26
27
-from vllm.attention.layer import Attention
27
from vllm.compilation.decorators import support_torch_compile
28
from vllm.config import CacheConfig, VllmConfig
29
from vllm.distributed import (
@@ -33,7 +32,8 @@ from vllm.distributed import (
32
tensor_model_parallel_all_gather,
33
)
34
from vllm.logger import init_logger
36
-from vllm.model_executor.layers.attention.chunked_local_attention import (
35
+from vllm.model_executor.layers.attention import (
36
+ Attention,
37
ChunkedLocalAttention,
38
)
39
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
vllm/model_executor/models/mimo_v2_flash.py
+1
-1
@@ -6,7 +6,6 @@ from itertools import islice
6
import torch
7
from torch import nn
8
9
-from vllm.attention.layer import Attention
9
from vllm.config import (
10
CacheConfig,
11
VllmConfig,
@@ -22,6 +21,7 @@ from vllm.distributed import (
21
)
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.activation import SiluAndMul
24
+from vllm.model_executor.layers.attention import Attention
25
from vllm.model_executor.layers.fused_moe import FusedMoE
26
from vllm.model_executor.layers.layernorm import RMSNorm
27
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/minicpm.py
+1
-1
@@ -33,7 +33,6 @@ import torch
33
from torch import nn
34
from transformers import PretrainedConfig
35
36
-from vllm.attention.layer import Attention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import (
@@ -43,6 +42,7 @@ from vllm.distributed import (
42
tensor_model_parallel_all_reduce,
43
)
44
from vllm.model_executor.layers.activation import FatreluAndMul, SiluAndMul
45
+from vllm.model_executor.layers.attention import Attention
46
from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/minicpm3.py
+1
-1
@@ -29,9 +29,9 @@ import torch
29
from torch import nn
30
from transformers import PretrainedConfig
31
32
-from vllm.attention.layer import Attention
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import get_tensor_model_parallel_world_size
34
+from vllm.model_executor.layers.attention import Attention
35
from vllm.model_executor.layers.layernorm import RMSNorm
36
from vllm.model_executor.layers.linear import (
37
ColumnParallelLinear,
vllm/model_executor/models/minimax_m2.py
+1
-1
@@ -30,7 +30,6 @@ import torch
30
from torch import nn
31
from transformers import PretrainedConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, ModelConfig, VllmConfig
35
from vllm.distributed import (
@@ -38,6 +37,7 @@ from vllm.distributed import (
37
get_tensor_model_parallel_world_size,
38
tensor_model_parallel_all_reduce,
39
)
40
+from vllm.model_executor.layers.attention import Attention
41
from vllm.model_executor.layers.fused_moe import FusedMoE
42
from vllm.model_executor.layers.layernorm import RMSNorm
43
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/minimax_text_01.py
+1
-1
@@ -14,7 +14,6 @@ import torch
14
from torch import nn
15
from transformers import MiniMaxConfig
16
17
-from vllm.attention.layer import Attention
17
from vllm.compilation.decorators import support_torch_compile
18
from vllm.config import CacheConfig, ModelConfig, VllmConfig
19
from vllm.distributed.parallel_state import (
@@ -24,6 +23,7 @@ from vllm.distributed.parallel_state import (
23
)
24
from vllm.forward_context import get_forward_context
25
from vllm.model_executor.layers.activation import SiluAndMul
26
+from vllm.model_executor.layers.attention import Attention
27
from vllm.model_executor.layers.fused_moe import FusedMoE
28
from vllm.model_executor.layers.layernorm import RMSNorm
29
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/mixtral.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
from torch import nn
33
from transformers import MixtralConfig
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
37
from vllm.distributed import (
@@ -40,6 +39,7 @@ from vllm.distributed import (
39
get_pp_group,
40
get_tensor_model_parallel_world_size,
41
)
42
+from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.fused_moe import FusedMoE
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/mllama4.py
+1
-1
@@ -36,7 +36,7 @@ from vllm.config import VllmConfig, set_current_vllm_config
36
from vllm.config.multimodal import BaseDummyOptions
37
from vllm.distributed import get_tensor_model_parallel_world_size
38
from vllm.forward_context import set_forward_context
39
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
39
+from vllm.model_executor.layers.attention import MMEncoderAttention
40
from vllm.model_executor.layers.fused_moe import FusedMoE
41
from vllm.model_executor.layers.linear import (
42
ColumnParallelLinear,
vllm/model_executor/models/modernbert.py
+1
-1
@@ -10,7 +10,7 @@ from transformers.activations import ACT2FN
10
from vllm.compilation.decorators import support_torch_compile
11
from vllm.config import ModelConfig, VllmConfig
12
from vllm.distributed import get_tensor_model_parallel_world_size
13
-from vllm.model_executor.layers.attention.encoder_only_attention import (
13
+from vllm.model_executor.layers.attention import (
14
EncoderOnlyAttention,
15
)
16
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
vllm/model_executor/models/molmo.py
+1
-2
@@ -17,7 +17,6 @@ from transformers import BatchFeature, PretrainedConfig, ProcessorMixin, TensorT
17
from transformers.image_utils import ImageInput
18
from transformers.tokenization_utils_base import TextInput
19
20
-from vllm.attention.layer import Attention
20
from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import CacheConfig, VllmConfig
22
from vllm.config.multimodal import BaseDummyOptions
@@ -29,7 +28,7 @@ from vllm.distributed import (
28
tensor_model_parallel_all_gather,
29
)
30
from vllm.model_executor.layers.activation import MulAndSilu, QuickGELU, SiluAndMul
32
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
31
+from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
32
from vllm.model_executor.layers.layernorm import RMSNorm
33
from vllm.model_executor.layers.linear import (
34
ColumnParallelLinear,
vllm/model_executor/models/molmo2.py
+1
-2
@@ -23,7 +23,6 @@ from transformers.image_utils import ImageInput
23
from transformers.tokenization_utils_base import TextInput
24
from transformers.video_utils import VideoInput, VideoMetadata
25
26
-from vllm.attention.layer import Attention
26
from vllm.compilation.decorators import support_torch_compile
27
from vllm.config import CacheConfig, VllmConfig
28
from vllm.config.multimodal import BaseDummyOptions, VideoDummyOptions
@@ -36,7 +35,7 @@ from vllm.distributed import (
35
)
36
from vllm.logger import init_logger
37
from vllm.model_executor.layers.activation import MulAndSilu, SiluAndMul, get_act_fn
39
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
38
+from vllm.model_executor.layers.attention import Attention, MMEncoderAttention
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
ColumnParallelLinear,
vllm/model_executor/models/moonvit.py
+1
-1
@@ -53,7 +53,7 @@ from transformers.activations import ACT2FN
53
from transformers.modeling_utils import PreTrainedModel
54
55
from vllm.distributed import divide, get_tensor_model_parallel_world_size
56
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
56
+from vllm.model_executor.layers.attention import MMEncoderAttention
57
from vllm.model_executor.layers.conv import Conv2dLayer
58
from vllm.model_executor.layers.linear import (
59
ColumnParallelLinear,
vllm/model_executor/models/mpt.py
+1
-1
@@ -10,7 +10,6 @@ import torch
10
import torch.nn as nn
11
from transformers import MptConfig
12
13
-from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import CacheConfig, VllmConfig
15
from vllm.distributed import (
@@ -19,6 +18,7 @@ from vllm.distributed import (
18
get_tensor_model_parallel_world_size,
19
)
20
from vllm.model_executor.layers.activation import get_act_fn
21
+from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.linear import (
23
ColumnParallelLinear,
24
QKVParallelLinear,
vllm/model_executor/models/nemotron.py
+1
-1
@@ -30,11 +30,11 @@ from itertools import islice
30
import torch
31
from torch import nn
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import get_act_fn
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.linear import (
39
ColumnParallelLinear,
40
QKVParallelLinear,
vllm/model_executor/models/nemotron_h.py
+1
-1
@@ -25,7 +25,6 @@ from itertools import islice
25
import torch
26
from torch import nn
27
28
-from vllm.attention.layer import Attention
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, ModelConfig, VllmConfig
30
from vllm.config.parallel import ParallelConfig
@@ -33,6 +32,7 @@ from vllm.distributed import get_ep_group, get_tensor_model_parallel_world_size
32
from vllm.distributed.communication_op import tensor_model_parallel_all_gather
33
from vllm.distributed.parallel_state import get_pp_group
34
from vllm.model_executor.layers.activation import ReLUSquaredActivation
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE
37
from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
38
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/olmo.py
+1
-1
@@ -31,11 +31,11 @@ import torch
31
from torch import nn
32
from transformers import OlmoConfig
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.linear import (
40
MergedColumnParallelLinear,
41
QKVParallelLinear,
vllm/model_executor/models/olmo2.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
from torch import nn
33
from transformers import Olmo2Config
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import VllmConfig
37
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -40,6 +39,7 @@ from vllm.distributed.communication_op import tensor_model_parallel_all_gather
39
from vllm.distributed.parallel_state import get_tensor_model_parallel_rank
40
from vllm.distributed.utils import split_tensor_along_last_dim
41
from vllm.model_executor.layers.activation import SiluAndMul
42
+from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.layernorm import RMSNorm
44
from vllm.model_executor.layers.linear import (
45
MergedColumnParallelLinear,
vllm/model_executor/models/olmoe.py
+1
-1
@@ -21,7 +21,6 @@ from itertools import islice
21
import torch
22
from torch import nn
23
24
-from vllm.attention.layer import Attention
24
from vllm.compilation.decorators import support_torch_compile
25
from vllm.config import VllmConfig
26
from vllm.distributed import (
@@ -32,6 +31,7 @@ from vllm.distributed import (
31
)
32
from vllm.distributed.utils import split_tensor_along_last_dim
33
from vllm.logger import init_logger
34
+from vllm.model_executor.layers.attention import Attention
35
from vllm.model_executor.layers.fused_moe import FusedMoE
36
from vllm.model_executor.layers.layernorm import RMSNorm
37
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/openpangu.py
+3
-2
@@ -29,7 +29,6 @@ import torch
29
from torch import nn
30
from transformers import PretrainedConfig
31
32
-from vllm.attention.layer import Attention, AttentionType
32
from vllm.compilation.decorators import support_torch_compile
33
from vllm.config import CacheConfig, ParallelConfig, VllmConfig
34
from vllm.distributed import (
@@ -41,7 +40,8 @@ from vllm.distributed import (
40
tensor_model_parallel_all_gather,
41
)
42
from vllm.model_executor.layers.activation import SiluAndMul
44
-from vllm.model_executor.layers.attention.static_sink_attention import (
43
+from vllm.model_executor.layers.attention import (
44
+ Attention,
45
StaticSinkAttention,
46
)
47
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
@@ -84,6 +84,7 @@ from vllm.model_executor.utils import set_weight_attrs
84
from vllm.platforms import current_platform
85
from vllm.sequence import IntermediateTensors
86
from vllm.transformers_utils.config import set_default_rope_theta
87
+from vllm.v1.attention.backend import AttentionType
88
from vllm.v1.attention.backends.flash_attn_diffkv import FlashAttentionDiffKVBackend
89
90
vllm/model_executor/models/opt.py
+1
-1
@@ -27,11 +27,11 @@ import torch
27
from torch import nn
28
from transformers import OPTConfig
29
30
-from vllm.attention.layer import Attention
30
from vllm.compilation.decorators import support_torch_compile
31
from vllm.config import CacheConfig, VllmConfig
32
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
33
from vllm.model_executor.layers.activation import get_act_fn
34
+from vllm.model_executor.layers.attention import Attention
35
from vllm.model_executor.layers.linear import (
36
ColumnParallelLinear,
37
QKVParallelLinear,
vllm/model_executor/models/orion.py
+1
-1
@@ -15,11 +15,11 @@ import torch
15
from torch import nn
16
from transformers import PretrainedConfig
17
18
-from vllm.attention.layer import Attention
18
from vllm.compilation.decorators import support_torch_compile
19
from vllm.config import CacheConfig, VllmConfig
20
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import SiluAndMul
22
+from vllm.model_executor.layers.attention import Attention
23
from vllm.model_executor.layers.linear import (
24
MergedColumnParallelLinear,
25
QKVParallelLinear,
vllm/model_executor/models/ouro.py
+1
-1
@@ -33,11 +33,11 @@ import torch
33
from torch import nn
34
from transformers import PretrainedConfig
35
36
-from vllm.attention.layer import Attention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_tensor_model_parallel_world_size
39
from vllm.model_executor.layers.activation import SiluAndMul
40
+from vllm.model_executor.layers.attention import Attention
41
from vllm.model_executor.layers.layernorm import RMSNorm
42
from vllm.model_executor.layers.linear import (
43
MergedColumnParallelLinear,
vllm/model_executor/models/paddleocr_vl.py
+1
-1
@@ -34,7 +34,7 @@ from vllm.config import VllmConfig
34
from vllm.config.multimodal import BaseDummyOptions
35
from vllm.distributed import parallel_state
36
from vllm.distributed import utils as dist_utils
37
-from vllm.model_executor.layers.attention.mm_encoder_attention import (
37
+from vllm.model_executor.layers.attention import (
38
MMEncoderAttention,
39
)
40
from vllm.model_executor.layers.conv import Conv2dLayer
vllm/model_executor/models/persimmon.py
+1
-1
@@ -30,11 +30,11 @@ import torch
30
from torch import nn
31
from transformers import PersimmonConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import get_act_fn
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.linear import (
39
ColumnParallelLinear,
40
QKVParallelLinear,
vllm/model_executor/models/phi.py
+1
-1
@@ -45,11 +45,11 @@ import torch
45
from torch import nn
46
from transformers import PhiConfig
47
48
-from vllm.attention.layer import Attention
48
from vllm.compilation.decorators import support_torch_compile
49
from vllm.config import CacheConfig, VllmConfig
50
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
51
from vllm.model_executor.layers.activation import get_act_fn
52
+from vllm.model_executor.layers.attention import Attention
53
from vllm.model_executor.layers.linear import (
54
ColumnParallelLinear,
55
QKVParallelLinear,
vllm/model_executor/models/phimoe.py
+1
-1
@@ -31,10 +31,10 @@ import torch
31
from torch import nn
32
from transformers.configuration_utils import PretrainedConfig
33
34
-from vllm.attention.layer import Attention
34
from vllm.compilation.decorators import support_torch_compile
35
from vllm.config import CacheConfig, VllmConfig
36
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.fused_moe import FusedMoE
39
from vllm.model_executor.layers.linear import (
40
QKVParallelLinear,
vllm/model_executor/models/plamo2.py
+1
-1
@@ -9,7 +9,6 @@ import torch
9
from torch import nn
10
from transformers import PretrainedConfig
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import VllmConfig, get_current_vllm_config
14
from vllm.distributed import divide, get_tensor_model_parallel_world_size
@@ -17,6 +16,7 @@ from vllm.distributed.parallel_state import get_pp_group
16
from vllm.forward_context import ForwardContext, get_forward_context
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.layers.activation import SiluAndMul
19
+from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.layernorm import RMSNorm
21
from vllm.model_executor.layers.linear import (
22
ColumnParallelLinear,
vllm/model_executor/models/plamo3.py
+1
-1
@@ -10,12 +10,12 @@ import torch
10
from torch import nn
11
from transformers import PretrainedConfig
12
13
-from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import VllmConfig
15
from vllm.distributed import get_tensor_model_parallel_world_size
16
from vllm.distributed.parallel_state import get_pp_group
17
from vllm.model_executor.layers.activation import SiluAndMul
18
+from vllm.model_executor.layers.attention import Attention
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import (
21
MergedColumnParallelLinear,
vllm/model_executor/models/qwen.py
+1
-1
@@ -16,11 +16,11 @@ import torch
16
from torch import nn
17
from transformers import PretrainedConfig
18
19
-from vllm.attention.layer import Attention
19
from vllm.compilation.decorators import support_torch_compile
20
from vllm.config import CacheConfig, VllmConfig
21
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
22
from vllm.model_executor.layers.activation import SiluAndMul
23
+from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
26
MergedColumnParallelLinear,
vllm/model_executor/models/qwen2.py
+2
-2
@@ -33,12 +33,12 @@ import torch
33
from torch import nn
34
from transformers import Qwen2Config
35
36
-from vllm.attention.layer import Attention
36
from vllm.compilation.decorators import support_torch_compile
37
from vllm.config import CacheConfig, VllmConfig
38
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
39
from vllm.model_executor.layers.activation import SiluAndMul
41
-from vllm.model_executor.layers.attention.encoder_only_attention import (
40
+from vllm.model_executor.layers.attention import (
41
+ Attention,
42
EncoderOnlyAttention,
43
)
44
from vllm.model_executor.layers.layernorm import RMSNorm
vllm/model_executor/models/qwen2_5_vl.py
+1
-1
@@ -49,7 +49,7 @@ from vllm.distributed import utils as dist_utils
49
from vllm.forward_context import set_forward_context
50
from vllm.logger import init_logger
51
from vllm.model_executor.layers.activation import get_act_and_mul_fn
52
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
52
+from vllm.model_executor.layers.attention import MMEncoderAttention
53
from vllm.model_executor.layers.conv import Conv3dLayer
54
from vllm.model_executor.layers.layernorm import RMSNorm
55
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/qwen2_moe.py
+1
-1
@@ -34,12 +34,12 @@ import torch.nn.functional as F
34
from torch import nn
35
from transformers import Qwen2MoeConfig
36
37
-from vllm.attention.layer import Attention
37
from vllm.compilation.decorators import support_torch_compile
38
from vllm.config import CacheConfig, VllmConfig
39
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
40
from vllm.logger import init_logger
41
from vllm.model_executor.layers.activation import SiluAndMul
42
+from vllm.model_executor.layers.attention import Attention
43
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
44
from vllm.model_executor.layers.layernorm import RMSNorm
45
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/qwen2_vl.py
+1
-1
@@ -49,7 +49,7 @@ from vllm.distributed import parallel_state, tensor_model_parallel_all_gather
49
from vllm.distributed import utils as dist_utils
50
from vllm.logger import init_logger
51
from vllm.model_executor.layers.activation import QuickGELU
52
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
52
+from vllm.model_executor.layers.attention import MMEncoderAttention
53
from vllm.model_executor.layers.conv import Conv3dLayer
54
from vllm.model_executor.layers.linear import (
55
ColumnParallelLinear,
vllm/model_executor/models/qwen3.py
+1
-1
@@ -30,11 +30,11 @@ import torch
30
from torch import nn
31
from transformers import Qwen3Config
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.logger import init_logger
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.layernorm import RMSNorm
39
from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
40
from vllm.model_executor.layers.logits_processor import LogitsProcessor
vllm/model_executor/models/qwen3_moe.py
+1
-1
@@ -32,7 +32,6 @@ import torch
32
import torch.nn.functional as F
33
from torch import nn
34
35
-from vllm.attention.layer import Attention
35
from vllm.compilation.decorators import support_torch_compile
36
from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
37
from vllm.distributed import (
@@ -43,6 +42,7 @@ from vllm.distributed import (
42
)
43
from vllm.logger import init_logger
44
from vllm.model_executor.layers.activation import SiluAndMul
45
+from vllm.model_executor.layers.attention import Attention
46
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
47
from vllm.model_executor.layers.layernorm import RMSNorm
48
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/qwen3_next.py
+1
-1
@@ -10,7 +10,6 @@ from einops import rearrange
10
from torch import nn
11
from transformers.activations import ACT2FN
12
13
-from vllm.attention.layer import Attention
13
from vllm.compilation.decorators import support_torch_compile
14
from vllm.config import (
15
CacheConfig,
@@ -29,6 +28,7 @@ from vllm.distributed import (
28
)
29
from vllm.forward_context import ForwardContext, get_forward_context
30
from vllm.logger import init_logger
31
+from vllm.model_executor.layers.attention import Attention
32
from vllm.model_executor.layers.fla.ops import (
33
chunk_gated_delta_rule,
34
fused_recurrent_gated_delta_rule,
vllm/model_executor/models/seed_oss.py
+1
-1
@@ -30,12 +30,12 @@ import torch
30
from torch import nn
31
from transformers import PretrainedConfig as SeedOssConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.logger import init_logger
37
from vllm.model_executor.layers.activation import SiluAndMul
38
+from vllm.model_executor.layers.attention import Attention
39
from vllm.model_executor.layers.layernorm import RMSNorm
40
from vllm.model_executor.layers.linear import (
41
MergedColumnParallelLinear,
vllm/model_executor/models/siglip.py
+2
-2
@@ -19,10 +19,10 @@ from vllm.config import VllmConfig
19
from vllm.config.multimodal import BaseDummyOptions
20
from vllm.distributed import divide, get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import get_act_fn
22
-from vllm.model_executor.layers.attention.encoder_only_attention import (
22
+from vllm.model_executor.layers.attention import (
23
EncoderOnlyAttention,
24
+ MMEncoderAttention,
25
)
25
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
26
from vllm.model_executor.layers.conv import Conv2dLayer
27
from vllm.model_executor.layers.linear import (
28
ColumnParallelLinear,
vllm/model_executor/models/siglip2navit.py
+1
-1
@@ -13,7 +13,7 @@ from transformers.configuration_utils import PretrainedConfig
13
14
from vllm.distributed import divide, get_tensor_model_parallel_world_size
15
from vllm.model_executor.layers.activation import get_act_fn
16
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
16
+from vllm.model_executor.layers.attention import MMEncoderAttention
17
from vllm.model_executor.layers.conv import Conv2dLayer
18
from vllm.model_executor.layers.linear import (
19
ColumnParallelLinear,
vllm/model_executor/models/solar.py
+1
-1
@@ -30,11 +30,11 @@ import torch
30
from torch import nn
31
from transformers import PretrainedConfig
32
33
-from vllm.attention.layer import Attention
33
from vllm.compilation.decorators import support_torch_compile
34
from vllm.config import CacheConfig, VllmConfig
35
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
36
from vllm.model_executor.layers.activation import SiluAndMul
37
+from vllm.model_executor.layers.attention import Attention
38
from vllm.model_executor.layers.layernorm import RMSNorm
39
from vllm.model_executor.layers.linear import (
40
MergedColumnParallelLinear,
vllm/model_executor/models/stablelm.py
+1
-1
@@ -29,10 +29,10 @@ import torch
29
from torch import nn
30
from transformers import StableLmConfig
31
32
-from vllm.attention.layer import Attention
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.model_executor.layers.activation import SiluAndMul
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.linear import (
37
MergedColumnParallelLinear,
38
QKVParallelLinear,
vllm/model_executor/models/starcoder2.py
+1
-1
@@ -28,11 +28,11 @@ import torch
28
from torch import nn
29
from transformers import Starcoder2Config
30
31
-from vllm.attention.layer import Attention
31
from vllm.compilation.decorators import support_torch_compile
32
from vllm.config import CacheConfig, VllmConfig
33
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
34
from vllm.model_executor.layers.activation import get_act_fn
35
+from vllm.model_executor.layers.attention import Attention
36
from vllm.model_executor.layers.linear import (
37
ColumnParallelLinear,
38
QKVParallelLinear,
vllm/model_executor/models/step1.py
+2
-1
@@ -10,7 +10,6 @@ from collections.abc import Iterable
10
import torch
11
from torch import nn
12
13
-from vllm.attention.layer import Attention, AttentionType
13
from vllm.config import CacheConfig, VllmConfig
14
from vllm.distributed import (
15
get_pp_group,
@@ -18,6 +17,7 @@ from vllm.distributed import (
17
get_tensor_model_parallel_world_size,
18
)
19
from vllm.model_executor.layers.activation import SiluAndMul
20
+from vllm.model_executor.layers.attention import Attention
21
from vllm.model_executor.layers.layernorm import RMSNorm
22
from vllm.model_executor.layers.linear import (
23
MergedColumnParallelLinear,
@@ -41,6 +41,7 @@ from vllm.model_executor.models.utils import (
41
maybe_prefix,
42
)
43
from vllm.sequence import IntermediateTensors
44
+from vllm.v1.attention.backend import AttentionType
45
46
STEP_PACKED_MODULES_MAPPING = {
47
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
vllm/model_executor/models/step3_text.py
+1
-1
@@ -9,7 +9,6 @@ from typing import Any
9
import torch
10
from torch import nn
11
12
-from vllm.attention.layer import Attention
12
from vllm.compilation.decorators import support_torch_compile
13
from vllm.config import CacheConfig, ModelConfig, VllmConfig
14
from vllm.distributed import (
@@ -19,6 +18,7 @@ from vllm.distributed import (
18
)
19
from vllm.logger import init_logger
20
from vllm.model_executor.layers.activation import SiluAndMul
21
+from vllm.model_executor.layers.attention import Attention
22
from vllm.model_executor.layers.fused_moe import FusedMoE
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (
vllm/model_executor/models/step3_vl.py
+1
-1
@@ -19,7 +19,7 @@ from vllm.config import VllmConfig
19
from vllm.config.multimodal import BaseDummyOptions
20
from vllm.distributed import get_tensor_model_parallel_world_size
21
from vllm.model_executor.layers.activation import get_act_fn
22
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
22
+from vllm.model_executor.layers.attention import MMEncoderAttention
23
from vllm.model_executor.layers.conv import Conv2dLayer
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
vllm/model_executor/models/transformers/base.py
+2
-2
@@ -27,12 +27,12 @@ from torch import nn
27
from transformers import AutoModel
28
from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
29
30
-from vllm.attention.layer import Attention
30
from vllm.config.utils import getattr_iter
31
from vllm.distributed import get_pp_group, get_tp_group
32
from vllm.distributed.utils import get_pp_indices
33
from vllm.logger import init_logger
35
-from vllm.model_executor.layers.attention.encoder_only_attention import (
34
+from vllm.model_executor.layers.attention import (
35
+ Attention,
36
EncoderOnlyAttention,
37
)
38
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
vllm/model_executor/models/whisper.py
+5
-3
@@ -17,7 +17,6 @@ from transformers import (
17
)
18
from transformers.models.whisper.modeling_whisper import sinusoids
19
20
-from vllm.attention.layer import Attention
20
from vllm.compilation.decorators import support_torch_compile
21
from vllm.config import CacheConfig, ModelConfig, SpeechToTextConfig, VllmConfig
22
from vllm.config.multimodal import BaseDummyOptions
@@ -25,8 +24,11 @@ from vllm.distributed import get_tensor_model_parallel_world_size
24
from vllm.inputs.data import PromptType
25
from vllm.logger import init_logger
26
from vllm.model_executor.layers.activation import get_act_fn
28
-from vllm.model_executor.layers.attention.cross_attention import CrossAttention
29
-from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
27
+from vllm.model_executor.layers.attention import (
28
+ Attention,
29
+ CrossAttention,
30
+ MMEncoderAttention,
31
+)
32
from vllm.model_executor.layers.linear import (
33
ColumnParallelLinear,
34
QKVParallelLinear,
vllm/model_executor/models/whisper_causal.py
+1
-1
@@ -10,9 +10,9 @@ import torch
10
import torch.nn.functional as F
11
from torch import nn
12
13
-from vllm.attention.layer import Attention
13
from vllm.config import CacheConfig, VllmConfig
14
from vllm.distributed import get_tensor_model_parallel_world_size
15
+from vllm.model_executor.layers.attention import Attention
16
from vllm.model_executor.layers.layernorm import RMSNorm
17
from vllm.model_executor.layers.linear import (
18
QKVParallelLinear,
vllm/model_executor/models/zamba2.py
+1
-1
@@ -16,11 +16,11 @@ import torch
16
from torch import nn
17
from transformers import Zamba2Config
18
19
-from vllm.attention.layer import Attention
19
from vllm.compilation.decorators import support_torch_compile
20
from vllm.config import CacheConfig, ModelConfig, VllmConfig
21
from vllm.distributed import get_tensor_model_parallel_world_size
22
from vllm.model_executor.layers.activation import GeluAndMul
23
+from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.layernorm import RMSNorm
25
from vllm.model_executor.layers.linear import (
26
ColumnParallelLinear,
vllm/v1/attention/backends/flash_attn.py
+1
-1
@@ -9,7 +9,7 @@ from typing import ClassVar
9
import numpy as np
10
import torch
11
12
-from vllm.attention.layer import Attention
12
+from vllm.model_executor.layers.attention import Attention
13
from vllm.v1.attention.backend import (
14
AttentionBackend,
15
AttentionImpl,
vllm/v1/attention/backends/rocm_aiter_fa.py
+1
-1
@@ -8,9 +8,9 @@ from typing import ClassVar
8
import torch
9
10
from vllm._aiter_ops import rocm_aiter_ops
11
-from vllm.attention.layer import Attention
11
from vllm.config import VllmConfig, get_layers_from_vllm_config
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.attention import Attention
14
from vllm.platforms import current_platform
15
from vllm.utils.math_utils import cdiv
16
from vllm.utils.platform_utils import get_cu_count
vllm/v1/spec_decode/draft_model.py
+1
-1
@@ -4,10 +4,10 @@ from typing import Any
4
5
import torch
6
7
-from vllm.attention.layer import Attention
7
from vllm.config import VllmConfig, get_layers_from_vllm_config
8
from vllm.config.speculative import SpeculativeConfig
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.attention import Attention
11
from vllm.model_executor.model_loader import get_model
12
from vllm.triton_utils import tl, triton
13
from vllm.v1.attention.backends.utils import (
vllm/v1/worker/gpu_model_runner.py
+1
-1
@@ -21,7 +21,6 @@ import torch.nn as nn
21
from tqdm import tqdm
22
23
import vllm.envs as envs
24
-from vllm.attention.layer import Attention, MLAAttention
24
from vllm.compilation.counter import compilation_counter
25
from vllm.compilation.cuda_graph import CUDAGraphStat, CUDAGraphWrapper
26
from vllm.compilation.monitor import set_cudagraph_capturing_enabled
@@ -50,6 +49,7 @@ from vllm.forward_context import (
49
)
50
from vllm.logger import init_logger
51
from vllm.lora.layers import LoRAMapping, LoRAMappingType
52
+from vllm.model_executor.layers.attention import Attention, MLAAttention
53
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
54
from vllm.model_executor.layers.fused_moe.routed_experts_capturer import (
55
RoutedExpertsCapturer,
vllm/v1/worker/utils.py
+1
-1
@@ -7,9 +7,9 @@ from dataclasses import dataclass, field
7
import torch
8
from typing_extensions import deprecated
9
10
-from vllm.attention.layer import Attention
10
from vllm.config import CacheConfig, VllmConfig
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention import Attention
13
from vllm.model_executor.models.interfaces import MultiModalEmbeddings
14
from vllm.model_executor.models.utils import extract_layer_index
15
from vllm.multimodal.registry import MultiModalRegistry