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

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


52
FP8_DTYPE = current_platform.fp8_dtype()
53
54
logger = init_logger(__name__)

55

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

64
        return is_flash_attn_2_available()
65
66
    if current_platform.is_rocm():
        from importlib.util import find_spec
67

68
        return find_spec("flash_attn") is not None
69
70
71
    return False


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

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

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

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

    return attn_backend, flash_attn_varlen_func


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
173
174
175
176
177
178
179
180
181
182
183
184
185
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)


186
class Attention(nn.Module, AttentionLayerBase):
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
    """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,
203
204
205
206
207
208
        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,
209
        prefix: str = "",
210
        attn_type: str = AttentionType.DECODER,
211
212
        kv_sharing_target_layer_name: str | None = None,
        attn_backend: type[AttentionBackend] | None = None,
213
        **extra_impl_args,
214
    ) -> None:
215
216
217
218
        """
        The KV cache is stored inside this class and is accessed via
        `self.kv_cache`.
        """
219
        super().__init__()
220
221
222
223
224
225
226
227
228
        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

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

247
248
249
250
        # Initialize KV cache quantization attributes
        _init_kv_cache_quant(
            self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
        )
251

252
253
254
255
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
        self.sliding_window = sliding_window
256
        self.has_sink = extra_impl_args.get("sinks") is not None
257

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

        impl_cls = self.attn_backend.get_impl_cls()
275
276
277
278
279
280
281
282
283
284
285
286
287
        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,
        )
288
289
        backend_name = self.attn_backend.get_name()
        self.backend = AttentionBackendEnum.__members__.get(backend_name)
290
        self.dtype = dtype
291

292
293
294
295
        # 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.
296
        self.use_direct_call = not current_platform.opaque_attention_op()
297

298
        self.use_output = self.attn_backend.accept_output_buffer
299
        compilation_config = vllm_config.compilation_config
300
301
302
303
        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
304
        self.attn_type = attn_type
305
306
307
308
309
310
311
312
313

        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

314
315
316
317
        # 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 = [
318
            torch.tensor([])
319
            for _ in range(vllm_config.parallel_config.pipeline_parallel_size)
320
        ]
321

322
323
324
325
        # 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)
326

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

335
336
337
338
339
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
340
341
342
        # 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.
343
        output_shape: torch.Size | None = None,
344
    ) -> torch.Tensor:
345
346
347
348
349
350
351
352
353
        """
        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
354
        if self.calculate_kv_scales:
355
            torch.ops.vllm.maybe_calc_kv_scales(query, key, value, self.layer_name)
356
357
358
359
360
361
362
363
        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"}
364
365
366
367

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

369
        if self.use_output:
370
            output_shape = output_shape if output_shape is not None else query.shape
371
            output = torch.empty(output_shape, dtype=output_dtype, device=query.device)
372
            hidden_size = output_shape[-1]
373
374
375
376
377
378
379
380
381
            # 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)
382
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
383
                forward_context: ForwardContext = get_forward_context()
384
                attn_metadata = forward_context.attn_metadata
385
386
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
387
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
388
389
390
                self.impl.forward(
                    self, query, key, value, self_kv_cache, attn_metadata, output=output
                )
391
392
            else:
                torch.ops.vllm.unified_attention_with_output(
393
394
                    query, key, value, output, self.layer_name
                )
395
            return output.view(-1, hidden_size)
396
        else:
397
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
398
                forward_context = get_forward_context()
399
                attn_metadata = forward_context.attn_metadata
400
401
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
402
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
403
404
405
                return self.impl.forward(
                    self, query, key, value, self_kv_cache, attn_metadata
                )
406
407
            else:
                return torch.ops.vllm.unified_attention(
408
409
                    query, key, value, self.layer_name
                )
410

411
412
    def calc_kv_scales(self, query, key, value):
        self._q_scale.copy_(torch.abs(query).max() / self.q_range)
413
414
        self._k_scale.copy_(torch.abs(key).max() / self.k_range)
        self._v_scale.copy_(torch.abs(value).max() / self.v_range)
415
        self._q_scale_float = self._q_scale.item()
416
417
418
419
420
        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

421
422
423
424
425
    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
426
        s += f", backend={self.impl.__class__.__name__}"
427
        return s
428

429
    def process_weights_after_loading(self, act_dtype: torch.dtype):
430
        self.impl.process_weights_after_loading(act_dtype)
431

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

435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
    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,
            )

459

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

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
468
        num_kv_heads: int | None = None,
469
470
471
        # This has no effect, it is only here to make it easier to swap
        # between Attention and MultiHeadAttention
        prefix: str = "",
472
        multimodal_config: MultiModalConfig | None = None,
473
    ) -> None:
474
475
476
477
478
        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
479
        self.layer_name = prefix
480

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

487
488
        # During model initialization, the default dtype is set as the model
        # weight and activation dtype.
489
        dtype = torch.get_default_dtype()
490
491

        # Determine the attention backend
492
493
494
495
496
497
498
499
        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,
        )
500
501
502
503
504
505

        # 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

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

518
519
        self.attn_backend, self._flash_attn_varlen_func = (
            maybe_get_vit_flash_attn_backend(
520
521
                self.attn_backend,
                use_upstream_fa,
522
                attn_backend_override=attn_backend_override,
523
            )
524
        )
525
526

        self.is_flash_attn_backend = self.attn_backend in {
527
528
            AttentionBackendEnum.FLASH_ATTN,
            AttentionBackendEnum.ROCM_AITER_FA,
529
530
531
532
        }

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

        logger.info_once(
            f"MultiHeadAttention attn_backend: {self.attn_backend}, "
541
542
            f"use_upstream_fa: {use_upstream_fa}"
        )
543

544
545
546
547
548
549
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> torch.Tensor:
550
        """Input shape:
551
552
553
554
        (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]
555
556
557
558
559
560
        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)

561
562
563
564
565
        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)

566
        if self.is_flash_attn_backend:
567
            assert self._flash_attn_varlen_func is not None
568
569
570
571
572
573
            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
            )
574
575
576
577
578
579
580
581
582
583
584

            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,
            )
585
        elif self.attn_backend == AttentionBackendEnum.TORCH_SDPA:
586
587
            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)
588
            out = out.transpose(1, 2)
589
        elif self.attn_backend == AttentionBackendEnum.PALLAS:
590
            query, key, value = (x.transpose(1, 2) for x in (query, key, value))
591
            from torch_xla.experimental.custom_kernel import flash_attention
592

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

601
        return out.reshape(bsz, q_len, -1)
602
603


604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
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,
622
        q_lora_rank: int | None,
623
624
        kv_lora_rank: int,
        kv_b_proj: ColumnParallelLinear,
625
626
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
627
628
        prefix: str = "",
        use_sparse: bool = False,
629
        indexer: object | None = None,
630
        **extra_impl_args,
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
    ):
        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
651
652
653
654
655

        # Initialize KV cache quantization attributes
        _init_kv_cache_quant(
            self, quant_config, prefix, kv_cache_dtype, calculate_kv_scales
        )
656
657
658
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

        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,
687
            **extra_impl_args,
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
        )

        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.
707
708
709
        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)
710
711
712
713
714
715

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

721
722
723
724
725
726
727
728
        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:
729
                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
                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:
746
                output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
747
748
749
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
                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

792
793
794
795
796
797
798
799
800
801
802
803
    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,
        )

804

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

814
815
816
    # 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:
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
        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,
)


839
def get_attention_context(
840
    layer_name: str,
841
842
843
844
845
846
847
848
) -> 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.
849

850
851
852
853
854
855
856
857
858
859
    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.
    """
860
    forward_context: ForwardContext = get_forward_context()
861
    attn_metadata = forward_context.attn_metadata
862
863
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
864
865
866
867
868
869
870
871
872
873
874
875
876
    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)
877
    output = self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
878
879

    return output
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895


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,
)
896
897


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


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


direct_register_custom_op(
    op_name="unified_attention_with_output",
    op_func=unified_attention_with_output,
937
    mutates_args=["output", "output_block_scale"],
938
939
    fake_impl=unified_attention_with_output_fake,
)
940
941


942
@maybe_transfer_kv_layer
943
944
945
946
947
948
def unified_mla_attention(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
949
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
    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,
)


973
@maybe_transfer_kv_layer
974
975
976
977
978
979
def unified_mla_attention_with_output(
    q: torch.Tensor,
    kv_c_normed: torch.Tensor,
    k_pe: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
980
981
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
982
) -> None:
983
    attn_metadata, self, kv_cache = get_attention_context(layer_name)
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
    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,
1003
1004
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
) -> 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,
)