layer.py 36 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer."""
4

5
6
from collections.abc import Callable
from typing import cast
7
8
9

import torch
import torch.nn as nn
10
import torch.nn.functional as F
11

12
import vllm.envs as envs
13
14
15
16
17
from vllm.attention.backends.abstract import (
    AttentionBackend,
    AttentionType,
    MLAAttentionImpl,
)
18
from vllm.attention.backends.registry import AttentionBackendEnum
19
from vllm.attention.selector import get_attn_backend
20
from vllm.attention.utils.kv_sharing_utils import validate_kv_sharing_target
21
from vllm.attention.utils.kv_transfer_utils import maybe_transfer_kv_layer
22
from vllm.config import CacheConfig, get_current_vllm_config
23
from vllm.config.multimodal import MultiModalConfig
24
from vllm.config.vllm import VllmConfig
25
from vllm.forward_context import ForwardContext, get_forward_context
26
from vllm.logger import init_logger
27
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
28
29
30
31
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    UnquantizedLinearMethod,
)
32
from vllm.model_executor.layers.quantization import QuantizationConfig
33
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
34
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
35
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
36
from vllm.model_executor.models.vision import get_vit_attn_backend
37
from vllm.platforms import current_platform
38
from vllm.utils.torch_utils import (
39
40
41
42
43
44
45
46
47
    direct_register_custom_op,
    kv_cache_dtype_str_to_dtype,
)
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheSpec,
    MLAAttentionSpec,
    SlidingWindowSpec,
)
48

49
50
51
52
53
54
if current_platform.is_rocm():
    from vllm.platforms.rocm import on_gfx9
else:
    on_gfx9 = lambda *args, **kwargs: False


55
FP8_DTYPE = current_platform.fp8_dtype()
56
57
logger = init_logger(__name__)

58

59
def check_upstream_fa_availability(dtype: torch.dtype):
60
61
62
63
64
    if (
        dtype in (torch.float16, torch.bfloat16)
        and current_platform.is_cuda()
        and current_platform.has_device_capability(80)
    ):
65
        from transformers.utils import is_flash_attn_2_available
66

67
        return is_flash_attn_2_available()
68
69
    if current_platform.is_rocm():
        from importlib.util import find_spec
70

71
        return find_spec("flash_attn") is not None
72
73
74
    return False


75
def maybe_get_vit_flash_attn_backend(
76
    attn_backend: AttentionBackendEnum,
77
    use_upstream_fa: bool,
78
79
    attn_backend_override: AttentionBackendEnum | None = None,
) -> tuple[AttentionBackendEnum, Callable | None]:
80
81
    if current_platform.is_rocm():
        if envs.VLLM_ROCM_USE_AITER and envs.VLLM_ROCM_USE_AITER_MHA and on_gfx9():
82
            attn_backend = AttentionBackendEnum.ROCM_AITER_FA
83
84
85
86
87
88

        elif (
            check_upstream_fa_availability(torch.get_default_dtype())
            and on_gfx9()
            and attn_backend_override is None
        ):
89
            attn_backend = AttentionBackendEnum.FLASH_ATTN
90
91
            use_upstream_fa = True
        else:
92
            return AttentionBackendEnum.TORCH_SDPA, None
93

94
    elif current_platform.is_cuda():
95
96
97
        if (
            attn_backend != AttentionBackendEnum.FLASH_ATTN
            and check_upstream_fa_availability(torch.get_default_dtype())
98
        ):
99
            attn_backend = AttentionBackendEnum.FLASH_ATTN
100
            use_upstream_fa = True
101
    elif current_platform.is_xpu():
102
        assert attn_backend == AttentionBackendEnum.FLASH_ATTN, (
103
104
105
            "XPU platform only supports FLASH_ATTN as vision attention backend."
        )
        use_upstream_fa = False
106
    else:
107
        return AttentionBackendEnum.TORCH_SDPA, None
108

109
110
111
112
113
    if attn_backend in {
        AttentionBackendEnum.FLASH_ATTN,
        AttentionBackendEnum.ROCM_AITER_FA,
    }:
        if attn_backend == AttentionBackendEnum.ROCM_AITER_FA:
114
115
116
117
118
            from aiter import flash_attn_varlen_func
        else:
            if use_upstream_fa:
                from flash_attn import flash_attn_varlen_func
            else:
119
                from vllm.attention.utils.fa_utils import flash_attn_varlen_func
120
121
122
123
124
125
    else:
        flash_attn_varlen_func = None

    return attn_backend, flash_attn_varlen_func


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
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
def _init_kv_cache_quant(
    layer: nn.Module,
    quant_config: QuantizationConfig | None,
    prefix: str,
    kv_cache_dtype: str,
    calculate_kv_scales: bool,
) -> None:
    """Initializes KV cache scaling factors and quantization method.

    This helper function sets up the KV cache quantization attributes that are
    shared between Attention and MLAAttention layers. It initializes scale
    tensors for query, key, value, and probability, and configures the
    quantization method if applicable.

    Args:
        layer: The attention layer instance to initialize.
        quant_config: Optional quantization configuration.
        prefix: Layer name prefix for quantization method lookup.
        kv_cache_dtype: The KV cache data type string.
        calculate_kv_scales: Whether to calculate KV scales dynamically.
    """
    # The default k/v_scale is set to 1.0. This is ignored
    # when kv-cache is not fp8, and should be used with
    # kv-cache in fp8_e5m2. For kv-cache in fp8_e4m3, we
    # expect the pre-quantized k/v_scale to be loaded along
    # with the model weights.
    layer.kv_cache_dtype = kv_cache_dtype
    layer.calculate_kv_scales = calculate_kv_scales
    layer._k_scale = torch.tensor(1.0, dtype=torch.float32)
    layer._v_scale = torch.tensor(1.0, dtype=torch.float32)
    layer._q_scale = torch.tensor(1.0, dtype=torch.float32)
    layer._prob_scale = torch.tensor(1.0, dtype=torch.float32)

    # We also keep q/k/v_scale on host (cpu) memory for attention
    # backends that require the scales to be on host instead of on device.
    # e.g. Flashinfer
    layer._q_scale_float = 1.0
    layer._k_scale_float = 1.0
    layer._v_scale_float = 1.0

    # The output scale on host memory. This should be the input scale of
    # the quant op after this attention layer.
    layer._o_scale_float = None

    quant_method = (
        quant_config.get_quant_method(layer, prefix=prefix) if quant_config else None
    )
    if quant_method is not None and not isinstance(
        quant_method, UnquantizedLinearMethod
    ):
        assert isinstance(quant_method, BaseKVCacheMethod)
        # TODO (mgoin): kv cache dtype should be specified in the FP8
        # checkpoint config and become the "auto" behavior
        if kv_cache_dtype == "fp8_e5m2":
            raise ValueError("fp8_e5m2 kv-cache is not supported with fp8 checkpoints.")
        # If quantization is enabled, we make "k_scale" and "v_scale"
        # parameters so that it can be loaded from the model checkpoint.
        # The k/v_scale will then be converted back to native float32
        # values after weight loading.
        layer.quant_method = quant_method
        layer.quant_method.create_weights(layer)


189
class Attention(nn.Module, AttentionLayerBase):
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
    """Attention layer.

    This class takes query, key, and value tensors as input. The input tensors
    can either contain prompt tokens or generation tokens.
    The class does the following:

    1. Store the input key and value tensors in the KV cache.
    2. Perform (multi-head/multi-query/grouped-query) attention.
    3. Return the output tensor.
    """

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
206
207
208
209
210
211
        num_kv_heads: int | None = None,
        alibi_slopes: list[float] | None = None,
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        logits_soft_cap: float | None = None,
        per_layer_sliding_window: int | None = None,
212
        prefix: str = "",
213
        attn_type: str = AttentionType.DECODER,
214
215
        kv_sharing_target_layer_name: str | None = None,
        attn_backend: type[AttentionBackend] | None = None,
216
        **extra_impl_args,
217
    ) -> None:
218
219
220
221
        """
        The KV cache is stored inside this class and is accessed via
        `self.kv_cache`.
        """
222
        super().__init__()
223
224
225
226
227
228
229
230
231
        if per_layer_sliding_window is not None:
            # per-layer sliding window
            sliding_window = per_layer_sliding_window
        elif cache_config is not None:
            # model-level sliding window
            sliding_window = cache_config.sliding_window
        else:
            sliding_window = None

232
        vllm_config = get_current_vllm_config()
233
234
235
        if cache_config is not None:
            kv_cache_dtype = cache_config.cache_dtype
            block_size = cache_config.block_size
236
            calculate_kv_scales = cache_config.calculate_kv_scales
237
238
239
        else:
            kv_cache_dtype = "auto"
            block_size = 16
240
            calculate_kv_scales = False
241
242
243
        self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
            kv_cache_dtype, vllm_config.model_config
        )
244
245
        if num_kv_heads is None:
            num_kv_heads = num_heads
246
247
248
        assert num_heads % num_kv_heads == 0, (
            f"num_heads ({num_heads}) is not divisible by num_kv_heads ({num_kv_heads})"
        )
249

250
251
252
253
        # Initialize KV cache quantization attributes
        _init_kv_cache_quant(
            self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
        )
254

255
256
257
258
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
        self.sliding_window = sliding_window
259
        self.has_sink = extra_impl_args.get("sinks") is not None
260

261
262
263
        # During model initialization, the default dtype is set as the model
        # weight and activation dtype.
        dtype = torch.get_default_dtype()
264
        if attn_backend is None:
265
266
267
268
269
            self.attn_backend = get_attn_backend(
                head_size,
                dtype,
                kv_cache_dtype,
                block_size,
270
                use_mla=False,
271
                has_sink=self.has_sink,
272
                attn_type=attn_type,
273
            )
274
275
276
277
        else:
            self.attn_backend = attn_backend

        impl_cls = self.attn_backend.get_impl_cls()
278
279
280
281
282
283
284
285
286
287
288
289
290
        self.impl = impl_cls(
            num_heads,
            head_size,
            scale,
            num_kv_heads,
            alibi_slopes,
            sliding_window,
            kv_cache_dtype,
            logits_soft_cap,
            attn_type,
            kv_sharing_target_layer_name,
            **extra_impl_args,
        )
291
292
        backend_name = self.attn_backend.get_name()
        self.backend = AttentionBackendEnum.__members__.get(backend_name)
293
        self.dtype = dtype
294

295
296
297
298
        # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
        # torch.compile works by registering the attention as one giant
        # opaque custom op. For other platforms, we directly call them
        # and let torch.compile handle them.
299
        self.use_direct_call = not current_platform.opaque_attention_op()
300

301
        self.use_output = self.attn_backend.accept_output_buffer
302
        compilation_config = vllm_config.compilation_config
303
304
305
306
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self
        self.layer_name = prefix
307
        self.attn_type = attn_type
308
309
310
311
312
313
314
315
316

        if kv_sharing_target_layer_name is not None:
            validate_kv_sharing_target(
                prefix,
                kv_sharing_target_layer_name,
                compilation_config.static_forward_context,
            )
        self.kv_sharing_target_layer_name = kv_sharing_target_layer_name

317
318
319
320
        # use a placeholder kv cache tensor during init, which will be replaced
        # by bind_kv_cache
        # this variable will not be accessed if use_direct_call is True
        self.kv_cache = [
321
            torch.tensor([])
322
            for _ in range(vllm_config.parallel_config.pipeline_parallel_size)
323
        ]
324

325
326
327
328
        # Initialize q/k/v range constants.
        self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
        self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
        self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
329

330
331
        # for attn backends supporting query quantization
        self.query_quant = None
332
333
        if (
            self.kv_cache_dtype.startswith("fp8")
334
            and self.impl.supports_quant_query_input()
335
336
        ):
            self.query_quant = QuantFP8(static=True, group_shape=GroupShape.PER_TENSOR)
337

338
339
340
341
342
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
343
344
345
        # For some alternate attention backends like MLA the attention output
        # shape does not match the query shape, so we optionally let the model
        # definition specify the output tensor shape.
346
        output_shape: torch.Size | None = None,
347
    ) -> torch.Tensor:
348
349
350
351
352
353
354
355
356
        """
        The KV cache is stored inside this class and is accessed via
        `self.kv_cache`.

        Attention metadata (`attn_metadata`) is set using a context manager in
        the model runner's `execute_model` method. It is accessed via forward
        context using
        `vllm.forward_context.get_forward_context().attn_metadata`.
        """
Chen Zhang's avatar
Chen Zhang committed
357
        if self.calculate_kv_scales:
358
            torch.ops.vllm.maybe_calc_kv_scales(query, key, value, self.layer_name)
359
360
361
362
363
364
365
366
        output_dtype = query.dtype
        if self.query_quant is not None:
            # quantizing with a simple torch operation enables
            # torch.compile to fuse this into previous ops
            # which reduces overheads during decoding.
            # Otherwise queries are quantized using custom ops
            # which causes decoding overheads
            assert self.kv_cache_dtype in {"fp8", "fp8_e4m3"}
367
368
369
370

            # check if query quantization is supported
            if self.impl.supports_quant_query_input():
                query, _ = self.query_quant(query, self._q_scale)
371

372
        if self.use_output:
373
            output_shape = output_shape if output_shape is not None else query.shape
374
            output = torch.empty(output_shape, dtype=output_dtype, device=query.device)
375
            hidden_size = output_shape[-1]
376
377
378
379
380
381
382
383
384
            # Reshape the query, key, and value tensors.
            # NOTE(woosuk): We do this outside the custom op to minimize the
            # CPU overheads from the non-CUDA-graph regions.
            query = query.view(-1, self.num_heads, self.head_size)
            output = output.view(-1, self.num_heads, self.head_size)
            if key is not None:
                key = key.view(-1, self.num_kv_heads, self.head_size)
            if value is not None:
                value = value.view(-1, self.num_kv_heads, self.head_size)
385
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
386
                forward_context: ForwardContext = get_forward_context()
387
                attn_metadata = forward_context.attn_metadata
388
389
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
390
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
391
392
393
                self.impl.forward(
                    self, query, key, value, self_kv_cache, attn_metadata, output=output
                )
394
395
            else:
                torch.ops.vllm.unified_attention_with_output(
396
397
                    query, key, value, output, self.layer_name
                )
398
            return output.view(-1, hidden_size)
399
        else:
400
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
401
                forward_context = get_forward_context()
402
                attn_metadata = forward_context.attn_metadata
403
404
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
405
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
406
407
408
                return self.impl.forward(
                    self, query, key, value, self_kv_cache, attn_metadata
                )
409
410
            else:
                return torch.ops.vllm.unified_attention(
411
412
                    query, key, value, self.layer_name
                )
413

414
415
    def calc_kv_scales(self, query, key, value):
        self._q_scale.copy_(torch.abs(query).max() / self.q_range)
416
417
        self._k_scale.copy_(torch.abs(key).max() / self.k_range)
        self._v_scale.copy_(torch.abs(value).max() / self.v_range)
418
        self._q_scale_float = self._q_scale.item()
419
420
421
422
423
        self._k_scale_float = self._k_scale.item()
        self._v_scale_float = self._v_scale.item()
        # We only calculate the scales once
        self.calculate_kv_scales = False

424
425
426
427
428
    def extra_repr(self) -> str:
        s = f"head_size={self.impl.head_size}"  # type: ignore
        s += f", num_heads={self.impl.num_heads}"  # type: ignore
        s += f", num_kv_heads={self.impl.num_kv_heads}"  # type: ignore
        s += f", scale={self.impl.scale}"  # type: ignore
429
        s += f", backend={self.impl.__class__.__name__}"
430
        return s
431

432
    def process_weights_after_loading(self, act_dtype: torch.dtype):
433
        self.impl.process_weights_after_loading(act_dtype)
434

435
436
437
    def get_attn_backend(self) -> type[AttentionBackend]:
        return self.attn_backend

438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        # Block size may get updated after model loading, refresh it
        block_size = vllm_config.cache_config.block_size
        # Should not be called for enc-dec or encoder-only attention.
        assert self.attn_type == AttentionType.DECODER
        if self.sliding_window is not None:
            assert not vllm_config.model_config.use_mla, (
                "MLA is not supported for slidingwindow"
            )
            return SlidingWindowSpec(
                block_size=block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                dtype=self.kv_cache_torch_dtype,
                sliding_window=self.sliding_window,
            )
        else:
            return FullAttentionSpec(
                block_size=block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                dtype=self.kv_cache_torch_dtype,
            )

462

463
464
465
466
467
468
469
470
class MultiHeadAttention(nn.Module):
    """Multi-headed attention without any cache, used for ViT."""

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
471
        num_kv_heads: int | None = None,
472
473
474
        # This has no effect, it is only here to make it easier to swap
        # between Attention and MultiHeadAttention
        prefix: str = "",
475
        multimodal_config: MultiModalConfig | None = None,
476
    ) -> None:
477
478
479
480
481
        super().__init__()
        self.num_heads = num_heads
        self.head_size = head_size
        self.scale = scale
        self.num_kv_heads = num_heads if num_kv_heads is None else num_kv_heads
482
        self.layer_name = prefix
483

484
485
        assert self.num_heads % self.num_kv_heads == 0, (
            f"num_heads ({self.num_heads}) is not "
486
            f"divisible by num_kv_heads ({self.num_kv_heads})"
487
        )
488
489
        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

490
491
        # During model initialization, the default dtype is set as the model
        # weight and activation dtype.
492
        dtype = torch.get_default_dtype()
493
494

        # Determine the attention backend
495
496
497
498
499
500
501
502
        attn_backend_override = None
        if multimodal_config is not None:
            attn_backend_override = multimodal_config.mm_encoder_attn_backend
        backend = get_vit_attn_backend(
            head_size=head_size,
            dtype=dtype,
            attn_backend_override=attn_backend_override,
        )
503
504
505
506
507
508

        # Some auto-selected backends can be upgraded
        # to upstream flash attention if available.
        # If vllm native fa is selected, we use it directly.
        use_upstream_fa = False

509
510
511
512
        self.attn_backend = (
            backend
            if backend
            in {
513
514
515
516
                AttentionBackendEnum.TORCH_SDPA,
                AttentionBackendEnum.PALLAS,
                AttentionBackendEnum.ROCM_AITER_FA,
                AttentionBackendEnum.FLASH_ATTN,
517
            }
518
            else AttentionBackendEnum.TORCH_SDPA
519
        )
520

521
522
        self.attn_backend, self._flash_attn_varlen_func = (
            maybe_get_vit_flash_attn_backend(
523
524
                self.attn_backend,
                use_upstream_fa,
525
                attn_backend_override=attn_backend_override,
526
            )
527
        )
528
529

        self.is_flash_attn_backend = self.attn_backend in {
530
531
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.ROCM_AITER_FA,
532
533
534
535
        }

        # this condition is just to make sure that the
        # use_upstream_fa in the log is correct
536
537
538
539
        if (
            current_platform.is_rocm()
            and self.attn_backend == AttentionBackendEnum.FLASH_ATTN
        ):
540
            use_upstream_fa = True
541
542
543

        logger.info_once(
            f"MultiHeadAttention attn_backend: {self.attn_backend}, "
544
545
            f"use_upstream_fa: {use_upstream_fa}"
        )
546

547
548
549
550
551
552
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> torch.Tensor:
553
        """Input shape:
554
555
556
557
        (batch_size x seq_len x hidden_size) or
        (batch_size x seq_len x num_heads x head_size)
        """
        bsz, q_len = query.size()[:2]
558
559
560
561
562
563
        kv_len = key.size(1)

        query = query.view(bsz, q_len, self.num_heads, self.head_size)
        key = key.view(bsz, kv_len, self.num_kv_heads, self.head_size)
        value = value.view(bsz, kv_len, self.num_kv_heads, self.head_size)

564
565
566
567
568
        if (num_repeat := self.num_queries_per_kv) > 1:
            # Handle MQA and GQA
            key = torch.repeat_interleave(key, num_repeat, dim=2)
            value = torch.repeat_interleave(value, num_repeat, dim=2)

569
        if self.is_flash_attn_backend:
570
            assert self._flash_attn_varlen_func is not None
571
572
573
574
575
576
            cu_seqlens_q = torch.arange(
                0, (bsz + 1) * q_len, step=q_len, dtype=torch.int32, device=query.device
            )
            cu_seqlens_k = torch.arange(
                0, (bsz + 1) * kv_len, step=kv_len, dtype=torch.int32, device=key.device
            )
577
578
579
580
581
582
583
584
585
586
587

            out = self._flash_attn_varlen_func(
                query.flatten(0, 1),
                key.flatten(0, 1),
                value.flatten(0, 1),
                cu_seqlens_q=cu_seqlens_q,
                cu_seqlens_k=cu_seqlens_k,
                max_seqlen_q=q_len,
                max_seqlen_k=kv_len,
                softmax_scale=self.scale,
            )
588
        elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
589
590
            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
            out = F.scaled_dot_product_attention(query, key, value, scale=self.scale)
591
            out = out.transpose(1, 2)
592
        elif self.attn_backend == AttentionBackendEnum.PALLAS:
593
            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
594
            from torch_xla.experimental.custom_kernel import flash_attention
595

596
597
            out = flash_attention(query, key, value, sm_scale=self.scale)
            out = out.transpose(1, 2)
598
599
600
        else:
            # ViT attention hasn't supported this backend yet
            raise NotImplementedError(
601
602
                f"ViT attention hasn't supported {self.attn_backend} backend yet."
            )
603

604
        return out.reshape(bsz, q_len, -1)
605
606


607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
class MLAAttention(nn.Module, AttentionLayerBase):
    """Multi-Head Latent Attention layer.

    This class takes query, and compressed key/value tensors as input.
    The class does the following:

    1. Store the input key and value tensors in the KV cache.
    2. Perform (multi-head/multi-query/grouped-query) attention.
    3. Return the output tensor.
    """

    def __init__(
        self,
        num_heads: int,
        scale: float,
        qk_nope_head_dim: int,
        qk_rope_head_dim: int,
        v_head_dim: int,
625
        q_lora_rank: int | None,
626
627
        kv_lora_rank: int,
        kv_b_proj: ColumnParallelLinear,
628
629
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
630
631
        prefix: str = "",
        use_sparse: bool = False,
632
        indexer: object | None = None,
633
        **extra_impl_args,
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
    ):
        super().__init__()
        self.num_heads = num_heads
        self.scale = scale
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.head_size = kv_lora_rank + qk_rope_head_dim
        self.layer_name = prefix

        if cache_config is not None:
            kv_cache_dtype = cache_config.cache_dtype
            block_size = cache_config.block_size
            calculate_kv_scales = cache_config.calculate_kv_scales
        else:
            kv_cache_dtype = "auto"
            block_size = 16
            calculate_kv_scales = False
654
655
656
657
658

        # Initialize KV cache quantization attributes
        _init_kv_cache_quant(
            self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
        )
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689

        dtype = torch.get_default_dtype()
        self.attn_backend = get_attn_backend(
            self.head_size,
            dtype,
            kv_cache_dtype,
            block_size,
            use_mla=True,
            use_sparse=use_sparse,
        )
        impl_cls = cast(type[MLAAttentionImpl], self.attn_backend.get_impl_cls())
        self.impl = impl_cls(
            num_heads=self.num_heads,
            head_size=self.head_size,
            scale=self.scale,
            num_kv_heads=1,
            alibi_slopes=None,
            sliding_window=None,
            kv_cache_dtype=self.kv_cache_dtype,
            logits_soft_cap=None,
            attn_type=AttentionType.DECODER,
            kv_sharing_target_layer_name=None,
            # MLA Args
            q_lora_rank=self.q_lora_rank,
            kv_lora_rank=self.kv_lora_rank,
            qk_nope_head_dim=self.qk_nope_head_dim,
            qk_rope_head_dim=self.qk_rope_head_dim,
            qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
            v_head_dim=self.v_head_dim,
            kv_b_proj=kv_b_proj,
            indexer=indexer,
690
            **extra_impl_args,
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
        )

        self.use_direct_call = not current_platform.opaque_attention_op()

        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

        self.kv_cache = [
            torch.tensor([])
            for _ in range(
                get_current_vllm_config().parallel_config.pipeline_parallel_size
            )
        ]

        self.use_sparse = use_sparse

        # Initialize q/k/v range constants.
710
711
712
        self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
        self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
        self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)
713
714
715
716
717
718

    def forward(
        self,
        q: torch.Tensor,
        kv_c_normed: torch.Tensor,
        k_pe: torch.Tensor,
719
        output_shape: torch.Size | None = None,
720
    ) -> torch.Tensor:
721
722
723
        if self.calculate_kv_scales:
            torch.ops.vllm.maybe_calc_kv_scales(q, kv_c_normed, k_pe, self.layer_name)

724
725
726
727
728
729
730
731
        if self.use_direct_call:
            forward_context: ForwardContext = get_forward_context()
            attn_metadata = forward_context.attn_metadata
            if isinstance(attn_metadata, dict):
                attn_metadata = attn_metadata[self.layer_name]
            self_kv_cache = self.kv_cache[forward_context.virtual_engine]

            if self.attn_backend.accept_output_buffer:
732
                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
                self.impl.forward(
                    self,
                    q,
                    kv_c_normed,
                    k_pe,
                    self_kv_cache,
                    attn_metadata,
                    output=output,
                )
                return output
            else:
                return self.impl.forward(
                    self, q, kv_c_normed, k_pe, self_kv_cache, attn_metadata
                )
        else:
            if self.attn_backend.accept_output_buffer:
749
                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
                torch.ops.vllm.unified_mla_attention_with_output(
                    q,
                    kv_c_normed,
                    k_pe,
                    output,
                    self.layer_name,
                )
                return output
            else:
                return torch.ops.vllm.unified_mla_attention(
                    q,
                    kv_c_normed,
                    k_pe,
                    self.layer_name,
                )

    def process_weights_after_loading(self, act_dtype: torch.dtype):
        if hasattr(self.impl, "process_weights_after_loading"):
            self.impl.process_weights_after_loading(act_dtype)

    def calc_kv_scales(
        self, q: torch.Tensor, kv_c_normed: torch.Tensor, k_pe: torch.Tensor
    ) -> None:
        """Optional scale calculation for MLA inputs.

        Mirrors Attention.calc_kv_scales. Not all MLA backends require this
        """
        # Use safe defaults if ranges are not present
        q_range = getattr(self, "q_range", torch.tensor(1.0))
        k_range = getattr(self, "k_range", torch.tensor(1.0))
        v_range = getattr(self, "v_range", torch.tensor(1.0))

        self._q_scale.copy_(torch.abs(q).max() / q_range)
        # kv_c_normed is the compressed KV representation; use it for k/v
        kv_abs_max = torch.abs(kv_c_normed).max()
        self._k_scale.copy_(kv_abs_max / k_range)
        self._v_scale.copy_(kv_abs_max / v_range)
        self._q_scale_float = self._q_scale.item()
        self._k_scale_float = self._k_scale.item()
        self._v_scale_float = self._v_scale.item()
        self.calculate_kv_scales = False

    def get_attn_backend(self) -> type[AttentionBackend]:
        return self.attn_backend

795
796
797
798
799
800
801
802
803
804
805
806
    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        kv_cache_dtype = kv_cache_dtype_str_to_dtype(
            self.kv_cache_dtype, vllm_config.model_config
        )
        return MLAAttentionSpec(
            block_size=vllm_config.cache_config.block_size,
            num_kv_heads=1,
            head_size=self.head_size,
            dtype=kv_cache_dtype,
            cache_dtype_str=vllm_config.cache_config.cache_dtype,
        )

807

808
809
810
811
812
813
814
def maybe_calc_kv_scales(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> None:
    forward_context: ForwardContext = get_forward_context()
815
    self = forward_context.no_compile_layers[layer_name]
816

817
818
819
    # Only calculate if the layer's calculate_kv_scales flag is True
    # This flag gets set to False after the first forward pass
    if not self.calculate_kv_scales:
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
        return

    self.calc_kv_scales(query, key, value)


def maybe_calc_kv_scales_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> None:
    return


direct_register_custom_op(
    op_name="maybe_calc_kv_scales",
    op_func=maybe_calc_kv_scales,
    mutates_args=["query", "key", "value"],
    fake_impl=maybe_calc_kv_scales_fake,
)


842
def get_attention_context(
843
    layer_name: str,
844
845
846
847
848
849
850
851
) -> tuple[dict | object | None, Attention | MLAAttention, torch.Tensor]:
    """Extract attention context for a given layer.

    This helper function extracts the attention metadata, attention layer
    instance, and KV cache tensor for a specific layer.

    Args:
        layer_name: The name/identifier of the attention layer.
852

853
854
855
856
857
858
859
860
861
862
    Returns:
        A tuple containing:
        - attn_metadata: Attention metadata for this specific layer, or None if
            no metadata available
        - attn_layer: The attention layer instance (Attention or MLAAttention)
        - kv_cache: The KV cache tensor for current virtual engine

        Note: attn_metadata may be None, but attn_layer and kv_cache are always
        extracted from the forward context.
    """
863
    forward_context: ForwardContext = get_forward_context()
864
    attn_metadata = forward_context.attn_metadata
865
866
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
867
868
869
870
871
872
873
874
875
876
877
878
879
    attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
    kv_cache = attn_layer.kv_cache[forward_context.virtual_engine]
    return attn_metadata, attn_layer, kv_cache


@maybe_transfer_kv_layer
def unified_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
880
    output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
881
882

    return output
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898


def unified_attention_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    return torch.empty_like(query).contiguous()


direct_register_custom_op(
    op_name="unified_attention",
    op_func=unified_attention,
    fake_impl=unified_attention_fake,
)
899
900


901
@maybe_transfer_kv_layer
902
903
904
905
906
907
def unified_attention_with_output(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
908
909
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
910
) -> None:
911
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
912
913
914
915
916
917
918
919
920
921
922
    self.impl.forward(
        self,
        query,
        key,
        value,
        kv_cache,
        attn_metadata,
        output=output,
        output_scale=output_scale,
        output_block_scale=output_block_scale,
    )
923
924
925
926
927
928
929
930


def unified_attention_with_output_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
931
932
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
933
934
935
936
937
938
939
) -> None:
    return


direct_register_custom_op(
    op_name="unified_attention_with_output",
    op_func=unified_attention_with_output,
940
    mutates_args=["output", "output_block_scale"],
941
942
    fake_impl=unified_attention_with_output_fake,
)
943
944


945
@maybe_transfer_kv_layer
946
947
948
949
950
951
def unified_mla_attention(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
952
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
    output = self.impl.forward(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata)

    return output


def unified_mla_attention_fake(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
    return torch.empty_like(q).contiguous()


direct_register_custom_op(
    op_name="unified_mla_attention",
    op_func=unified_mla_attention,
    mutates_args=[],
    fake_impl=unified_mla_attention_fake,
    dispatch_key=current_platform.dispatch_key,
)


976
@maybe_transfer_kv_layer
977
978
979
980
981
982
def unified_mla_attention_with_output(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
983
984
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
985
) -> None:
986
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
    self.impl.forward(
        self,
        q,
        kv_c_normed,
        k_pe,
        kv_cache,
        attn_metadata,
        output=output,
        output_scale=output_scale,
        output_block_scale=output_block_scale,
    )


def unified_mla_attention_with_output_fake(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
1006
1007
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
) -> None:
    return


direct_register_custom_op(
    op_name="unified_mla_attention_with_output",
    op_func=unified_mla_attention_with_output,
    mutates_args=["output", "output_block_scale"],
    fake_impl=unified_mla_attention_with_output_fake,
    dispatch_key=current_platform.dispatch_key,
)