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 10, 05:15 UTC
[4/N][Attention] Move MLA common to model_executor (#32060)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Matthew Bonanni committed
Jan 13, 2026 at 12:08 UTC
2263d44b688902aa3fd384ecdd7e3db3460b01e0
14 files changed
+50
-44
tests/v1/attention/test_mla_backends.py
+1
-1
@@ -19,12 +19,12 @@ from tests.v1.attention.utils import (
19
)
20
from vllm import _custom_ops as ops
21
from vllm.config.vllm import set_current_vllm_config
22
+from vllm.model_executor.layers.attention.mla_attention import QueryLenSupport
23
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
24
from vllm.utils.math_utils import cdiv
25
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
26
from vllm.v1.attention.backend import CommonAttentionMetadata
27
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
27
-from vllm.v1.attention.backends.mla.common import QueryLenSupport
28
from vllm.v1.attention.backends.registry import AttentionBackendEnum
29
from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported
30
from vllm.v1.kv_cache_interface import FullAttentionSpec
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+1
-1
@@ -14,9 +14,9 @@ from vllm.distributed.kv_transfer.kv_connector.v1.base import (
14
KVConnectorRole,
15
)
16
from vllm.logger import init_logger
17
+from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
18
from vllm.utils.hashing import safe_hash
19
from vllm.v1.attention.backend import AttentionMetadata
19
-from vllm.v1.attention.backends.mla.common import MLACommonMetadata
20
from vllm.v1.core.sched.output import SchedulerOutput
21
22
if TYPE_CHECKING:
vllm/distributed/kv_transfer/kv_connector/v1/p2p/p2p_nccl_connector.py
+1
-1
@@ -18,8 +18,8 @@ from vllm.distributed.kv_transfer.kv_connector.v1.p2p.p2p_nccl_engine import (
18
)
19
from vllm.distributed.parallel_state import get_world_group
20
from vllm.logger import init_logger
21
+from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
22
from vllm.v1.attention.backend import AttentionMetadata
22
-from vllm.v1.attention.backends.mla.common import MLACommonMetadata
23
from vllm.v1.core.sched.output import SchedulerOutput
24
25
if TYPE_CHECKING:
vllm/logging_utils/formatter.py
-2
@@ -33,8 +33,6 @@ class NewLineFormatter(logging.Formatter):
33
model_executor/.../quantization/utils/fp8_utils.py
34
vllm/model_executor/layers/quantization/awq.py ->
35
model_executor/layers/quantization/awq.py
36
- vllm/v1/attention/backends/mla/common.py ->
37
- v1/attention/backends/mla/common.py
36
37
Args:
38
relpath (Path): The relative path to be shortened.
vllm/model_executor/layers/attention/mla_attention.py
renamed
vllm/v1/attention/backends/mla/cutlass_mla.py
+6
-6
@@ -9,6 +9,12 @@ import torch
9
import vllm._custom_ops as ops
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention.mla_attention import (
13
+ MLACommonBackend,
14
+ MLACommonImpl,
15
+ MLACommonMetadata,
16
+ MLACommonMetadataBuilder,
17
+)
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.v1.attention.backend import (
20
AttentionCGSupport,
@@ -17,12 +23,6 @@ from vllm.v1.attention.backend import (
23
MultipleOf,
24
is_quantized_kv_cache,
25
)
20
-from vllm.v1.attention.backends.mla.common import (
21
- MLACommonBackend,
22
- MLACommonImpl,
23
- MLACommonMetadata,
24
- MLACommonMetadataBuilder,
25
-)
26
27
logger = init_logger(__name__)
28
vllm/v1/attention/backends/mla/flashattn_mla.py
+8
-8
@@ -9,6 +9,14 @@ import torch
9
from vllm.config import VllmConfig
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention.mla_attention import (
13
+ MLACommonBackend,
14
+ MLACommonDecodeMetadata,
15
+ MLACommonImpl,
16
+ MLACommonMetadata,
17
+ MLACommonMetadataBuilder,
18
+ QueryLenSupport,
19
+)
20
from vllm.model_executor.layers.batch_invariant import (
21
vllm_is_batch_invariant,
22
)
@@ -24,14 +32,6 @@ from vllm.v1.attention.backends.fa_utils import (
32
flash_attn_supports_mla,
33
get_flash_attn_version,
34
)
27
-from vllm.v1.attention.backends.mla.common import (
28
- MLACommonBackend,
29
- MLACommonDecodeMetadata,
30
- MLACommonImpl,
31
- MLACommonMetadata,
32
- MLACommonMetadataBuilder,
33
- QueryLenSupport,
34
-)
35
from vllm.v1.kv_cache_interface import AttentionSpec
36
from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]
37
flash_attn_varlen_func,
vllm/v1/attention/backends/mla/flashinfer_mla.py
+7
-7
@@ -8,6 +8,13 @@ from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla
8
9
from vllm.config.cache import CacheDType
10
from vllm.logger import init_logger
11
+from vllm.model_executor.layers.attention.mla_attention import (
12
+ MLACommonBackend,
13
+ MLACommonImpl,
14
+ MLACommonMetadata,
15
+ MLACommonMetadataBuilder,
16
+ QueryLenSupport,
17
+)
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.v1.attention.backend import (
20
AttentionCGSupport,
@@ -15,13 +22,6 @@ from vllm.v1.attention.backend import (
22
AttentionType,
23
MultipleOf,
24
)
18
-from vllm.v1.attention.backends.mla.common import (
19
- MLACommonBackend,
20
- MLACommonImpl,
21
- MLACommonMetadata,
22
- MLACommonMetadataBuilder,
23
- QueryLenSupport,
24
-)
25
from vllm.v1.attention.backends.utils import KVCacheLayoutType
26
27
logger = init_logger(__name__)
vllm/v1/attention/backends/mla/flashmla.py
+8
-8
@@ -9,6 +9,14 @@ import torch
9
from vllm.config import VllmConfig
10
from vllm.config.cache import CacheDType
11
from vllm.logger import init_logger
12
+from vllm.model_executor.layers.attention.mla_attention import (
13
+ MLACommonBackend,
14
+ MLACommonDecodeMetadata,
15
+ MLACommonImpl,
16
+ MLACommonMetadata,
17
+ MLACommonMetadataBuilder,
18
+ QueryLenSupport,
19
+)
20
from vllm.model_executor.layers.batch_invariant import (
21
vllm_is_batch_invariant,
22
)
@@ -19,14 +27,6 @@ from vllm.v1.attention.backend import (
27
AttentionType,
28
MultipleOf,
29
)
22
-from vllm.v1.attention.backends.mla.common import (
23
- MLACommonBackend,
24
- MLACommonDecodeMetadata,
25
- MLACommonImpl,
26
- MLACommonMetadata,
27
- MLACommonMetadataBuilder,
28
- QueryLenSupport,
29
-)
30
from vllm.v1.attention.backends.utils import (
31
reshape_attn_output_for_spec_decode,
32
reshape_query_for_spec_decode,
vllm/v1/attention/backends/mla/flashmla_sparse.py
+4
-1
@@ -10,6 +10,10 @@ from vllm import _custom_ops as ops
10
from vllm.config import VllmConfig, get_current_vllm_config
11
from vllm.config.cache import CacheDType
12
from vllm.logger import init_logger
13
+from vllm.model_executor.layers.attention.mla_attention import (
14
+ MLACommonBaseImpl,
15
+ get_mla_dims,
16
+)
17
from vllm.platforms import current_platform
18
from vllm.platforms.interface import DeviceCapability
19
from vllm.triton_utils import tl, triton
@@ -23,7 +27,6 @@ from vllm.v1.attention.backend import (
27
CommonAttentionMetadata,
28
MultipleOf,
29
)
26
-from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
30
from vllm.v1.attention.backends.utils import (
31
reshape_attn_output_for_spec_decode,
32
reshape_query_for_spec_decode,
vllm/v1/attention/backends/mla/rocm_aiter_mla.py
+2
-2
@@ -8,8 +8,7 @@ import torch
8
9
from vllm._aiter_ops import rocm_aiter_ops
10
from vllm.config import VllmConfig
11
-from vllm.v1.attention.backend import AttentionCGSupport, AttentionLayer, MultipleOf
12
-from vllm.v1.attention.backends.mla.common import (
11
+from vllm.model_executor.layers.attention.mla_attention import (
12
MLACommonBackend,
13
MLACommonDecodeMetadata,
14
MLACommonImpl,
@@ -17,6 +16,7 @@ from vllm.v1.attention.backends.mla.common import (
16
MLACommonMetadataBuilder,
17
QueryLenSupport,
18
)
19
+from vllm.v1.attention.backend import AttentionCGSupport, AttentionLayer, MultipleOf
20
from vllm.v1.kv_cache_interface import AttentionSpec
21
22
vllm/v1/attention/backends/mla/rocm_aiter_mla_sparse.py
+4
-1
@@ -11,6 +11,10 @@ from vllm import _custom_ops as ops
11
from vllm._aiter_ops import rocm_aiter_ops
12
from vllm.config import VllmConfig
13
from vllm.logger import init_logger
14
+from vllm.model_executor.layers.attention.mla_attention import (
15
+ MLACommonBaseImpl,
16
+ get_mla_dims,
17
+)
18
from vllm.v1.attention.backend import (
19
AttentionBackend,
20
AttentionCGSupport,
@@ -19,7 +23,6 @@ from vllm.v1.attention.backend import (
23
AttentionMetadataBuilder,
24
CommonAttentionMetadata,
25
)
22
-from vllm.v1.attention.backends.mla.common import MLACommonBaseImpl, get_mla_dims
26
from vllm.v1.attention.backends.mla.flashmla_sparse import (
27
triton_convert_req_index_to_global_index,
28
)
vllm/v1/attention/backends/mla/triton_mla.py
+5
-5
@@ -7,6 +7,11 @@ import torch
7
8
from vllm.config.cache import CacheDType
9
from vllm.logger import init_logger
10
+from vllm.model_executor.layers.attention.mla_attention import (
11
+ MLACommonBackend,
12
+ MLACommonImpl,
13
+ MLACommonMetadata,
14
+)
15
from vllm.model_executor.layers.batch_invariant import (
16
vllm_is_batch_invariant,
17
)
@@ -16,11 +21,6 @@ from vllm.v1.attention.backend import (
21
AttentionType,
22
is_quantized_kv_cache,
23
)
19
-from vllm.v1.attention.backends.mla.common import (
20
- MLACommonBackend,
21
- MLACommonImpl,
22
- MLACommonMetadata,
23
-)
24
from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd
25
26
logger = init_logger(__name__)
vllm/v1/spec_decode/eagle.py
+3
-1
@@ -183,7 +183,9 @@ class EagleProposer:
183
rocm_types.append(AiterFlashAttentionMetadata)
184
185
# TRITON_MLA backend support for MLA models (e.g., DeepSeek)
186
- from vllm.v1.attention.backends.mla.common import MLACommonMetadata
186
+ from vllm.model_executor.layers.attention.mla_attention import (
187
+ MLACommonMetadata,
188
+ )
189
190
rocm_types.append(MLACommonMetadata)
191