Skip to content

vllm.v1.attention.backends.mla.prefill.selector

Selector for MLA prefill backends.

This module provides functions for selecting the appropriate MLA prefill backend based on device capabilities and configuration.

Classes:

Functions:

MLAPrefillSelectorConfig

Bases: NamedTuple

Hashable configuration for MLA prefill backend selection.

This is analogous to AttentionSelectorConfig and contains model-specific configuration needed to select an MLA prefill backend, extracted from VllmConfig into a hashable form for caching.

Source code in vllm/v1/attention/backends/mla/prefill/selector.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
class MLAPrefillSelectorConfig(NamedTuple):
    """Hashable configuration for MLA prefill backend selection.

    This is analogous to AttentionSelectorConfig and contains model-specific
    configuration needed to select an MLA prefill backend, extracted from
    VllmConfig into a hashable form for caching.
    """

    dtype: torch.dtype
    mla_dimensions: MLADimensions = MLADimensions(
        qk_nope_head_dim=0,
        qk_rope_head_dim=0,
        v_head_dim=0,
    )

    def __repr__(self):
        return (
            f"MLAPrefillSelectorConfig(dtype={self.dtype}, "
            f"mla_dimensions={self.mla_dimensions})"
        )

_auto_select_mla_prefill_backend(device_capability, selector_config) cached

Auto-select the best available MLA prefill backend.

Parameters:

Returns:

Source code in vllm/v1/attention/backends/mla/prefill/selector.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
@cache
def _auto_select_mla_prefill_backend(
    device_capability: DeviceCapability,
    selector_config: MLAPrefillSelectorConfig,
) -> "type[MLAPrefillBackend]":
    """Auto-select the best available MLA prefill backend.

    Args:
        device_capability: The device's compute capability.
        selector_config: Hashable configuration for backend selection.

    Returns:
        The selected prefill backend class.

    """
    priorities = _get_mla_prefill_backend_priorities(
        device_capability,
        selector_config.mla_dimensions,
    )
    all_invalid_reasons: dict[str, list[str]] = {}

    for backend_enum in priorities:
        backend_cls: type[MLAPrefillBackend] | None = None
        try:
            backend_cls = backend_enum.get_class()
            invalid_reasons = backend_cls.validate_configuration(
                device_capability, selector_config
            )
        except ImportError:
            invalid_reasons = ["ImportError"]
        if not invalid_reasons:
            assert backend_cls is not None
            logger.info_once("Using %s MLA prefill backend.", backend_enum.name)
            return backend_cls
        all_invalid_reasons[backend_enum.name] = invalid_reasons

    reasons_str = (
        "{"
        + ", ".join(
            f"{name}: [{', '.join(reasons)}]"
            for name, reasons in all_invalid_reasons.items()
        )
        + "}"
    )
    config_str = repr(selector_config)
    logger.debug_once(
        "Some MLA prefill backends are not valid with %s. Reasons: %s.",
        config_str,
        reasons_str,
    )

    raise ValueError(
        f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}."
    )

_get_mla_prefill_backend_priorities(device_capability, mla_dimensions)

Get MLA prefill backend priorities based on device capability.

Parameters:

  • device_capability

    (DeviceCapability) –

    The device's compute capability.

  • mla_dimensions

    (MLADimensions) –

    The model's MLA head dimensions.

Returns:

Source code in vllm/v1/attention/backends/mla/prefill/selector.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def _get_mla_prefill_backend_priorities(
    device_capability: DeviceCapability,
    mla_dimensions: MLADimensions,
) -> list[MLAPrefillBackendEnum]:
    """Get MLA prefill backend priorities based on device capability.

    Args:
        device_capability: The device's compute capability.
        mla_dimensions: The model's MLA head dimensions.

    Returns:
        List of backends in priority order (highest priority first).

    """
    from vllm.platforms import current_platform

    if current_platform.is_rocm():
        return [
            MLAPrefillBackendEnum.ROCM_AITER_FA,
            MLAPrefillBackendEnum.FLASH_ATTN,
        ]

    if device_capability.major == 10:  # Blackwell
        if mla_dimensions == MLADimensions(
            qk_nope_head_dim=192,
            qk_rope_head_dim=64,
            v_head_dim=256,
        ):
            return [
                MLAPrefillBackendEnum.TRTLLM_RAGGED,
                MLAPrefillBackendEnum.FLASH_ATTN,
                MLAPrefillBackendEnum.FLASHINFER,
                MLAPrefillBackendEnum.TOKENSPEED_MLA,
            ]
        return [
            MLAPrefillBackendEnum.FLASH_ATTN,
            MLAPrefillBackendEnum.TRTLLM_RAGGED,
            MLAPrefillBackendEnum.FLASHINFER,
            MLAPrefillBackendEnum.TOKENSPEED_MLA,
        ]
    else:  # Hopper (SM90) and older
        return [
            MLAPrefillBackendEnum.FLASH_ATTN,
        ]

get_mla_prefill_backend(vllm_config)

Select the MLA prefill backend based on configuration and device.

This function first checks for explicit user preferences via mla_prefill_backend in AttentionConfig, then falls back to automatic priority-based selection.

Parameters:

Returns:

Source code in vllm/v1/attention/backends/mla/prefill/selector.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def get_mla_prefill_backend(
    vllm_config: "VllmConfig",
) -> "type[MLAPrefillBackend]":
    """Select the MLA prefill backend based on configuration and device.

    This function first checks for explicit user preferences via
    mla_prefill_backend in AttentionConfig, then falls back to automatic
    priority-based selection.

    Args:
        vllm_config: The vLLM configuration.

    Returns:
        The selected prefill backend class.

    """
    from vllm.platforms import current_platform

    if current_platform.is_cpu():
        # CPUs have no compute capability, so the capability-driven priority
        # path below does not apply. Prefer an accelerator-specific CPU backend
        # when its kernels are present, else the generic SDPA one.
        for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU):
            try:
                cpu_backend_cls = backend_enum.get_class()
            except ImportError:
                continue
            if cpu_backend_cls.is_available():
                logger.info_once("Using %s MLA prefill backend.", backend_enum.name)
                return cpu_backend_cls
        raise ValueError("No valid CPU MLA prefill backend found.")

    device_capability = current_platform.get_device_capability()
    if device_capability is None:
        logger.info_once(
            "Device capability not available, using FlashAttention MLA prefill backend."
        )
        return MLAPrefillBackendEnum.FLASH_ATTN.get_class()

    attention_config = vllm_config.attention_config

    model_config = vllm_config.model_config
    if model_config is None:
        selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype())
    else:
        hf_text_config = model_config.hf_text_config
        selector_config = MLAPrefillSelectorConfig(
            dtype=model_config.dtype,
            mla_dimensions=MLADimensions(
                qk_nope_head_dim=getattr(hf_text_config, "qk_nope_head_dim", 0),
                qk_rope_head_dim=getattr(hf_text_config, "qk_rope_head_dim", 0),
                v_head_dim=getattr(hf_text_config, "v_head_dim", 0),
            ),
        )

    if attention_config.mla_prefill_backend is not None:
        selected_backend = attention_config.mla_prefill_backend
        backend_cls: type[MLAPrefillBackend] | None = None
        try:
            backend_cls = selected_backend.get_class()
            invalid_reasons = backend_cls.validate_configuration(
                device_capability, selector_config
            )
        except ImportError:
            invalid_reasons = ["ImportError"]
        if invalid_reasons:
            raise ValueError(
                f"Selected MLA prefill backend {selected_backend.name} "
                f"is not valid for this configuration. "
                f"Reason: {invalid_reasons}"
            )
        assert backend_cls is not None
        logger.info_once("Using %s MLA prefill backend.", selected_backend.name)
        return backend_cls

    return _auto_select_mla_prefill_backend(
        device_capability,
        selector_config,
    )