attention.py 30.8 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from typing import TYPE_CHECKING, Any
5
6
7
8

import torch
import torch.nn as nn

9
import vllm.envs as envs
10
from vllm.config import CacheConfig, get_current_vllm_config
11
from vllm.config.vllm import VllmConfig
12
from vllm.forward_context import ForwardContext, get_forward_context
13
from vllm.logger import init_logger
14
15
16
from vllm.model_executor.layers.attention.kv_transfer_utils import (
    maybe_transfer_kv_layer,
)
17
from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
18
19
20
from vllm.model_executor.layers.linear import (
    UnquantizedLinearMethod,
)
21
from vllm.model_executor.layers.quantization import QuantizationConfig
22
from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
23
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
24
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
25
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
26
from vllm.platforms import current_platform
27
from vllm.utils.torch_utils import (
28
29
30
    LayerNameType,
    _encode_layer_name,
    _resolve_layer_name,
31
32
33
    direct_register_custom_op,
    kv_cache_dtype_str_to_dtype,
)
34
35
from vllm.v1.attention.backend import (
    AttentionBackend,
36
    AttentionMetadata,
37
38
39
40
    AttentionType,
)
from vllm.v1.attention.backends.registry import AttentionBackendEnum
from vllm.v1.attention.selector import get_attn_backend
41
42
43
44
from vllm.v1.kv_cache_interface import (
    FullAttentionSpec,
    KVCacheSpec,
    SlidingWindowSpec,
45
    get_kv_quant_mode,
46
)
47

48
49
50
if TYPE_CHECKING:
    from vllm.model_executor.layers.attention import MLAAttention

51
52
logger = init_logger(__name__)

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
def validate_kv_sharing_target(
    current_layer_name, target_layer_name, static_forward_context
):
    error_msg = (
        f"Specified KV sharing target layer for {current_layer_name} "
        f"is not valid: target layer {target_layer_name} "
    )

    if current_layer_name == target_layer_name:
        raise ValueError(error_msg + "cannot be the same as the current layer.")

    if target_layer_name not in static_forward_context:
        from vllm.model_executor.models.utils import extract_layer_index

        # If target layer name is not in the static fwd context, it means either
        # a) the target layer does not come BEFORE the current layer, or
        # b) the target layer is not an Attention layer that exists in the model
        current_layer_idx = extract_layer_index(current_layer_name)
        target_layer_idx = extract_layer_index(target_layer_name)
        if current_layer_idx <= target_layer_idx:
            raise ValueError(error_msg + "must come before the current layer.")
        else:
            raise ValueError(error_msg + "is not a valid Attention layer in the model.")

    # Currently KV sharing is only supported between layers of the same type
    target_layer_attn_type = static_forward_context[target_layer_name].attn_type
    expected = static_forward_context[current_layer_name].attn_type
    if target_layer_attn_type != expected:
        raise ValueError(
            error_msg + f"must be the same type as the current layer ({expected})."
        )


87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def should_load_quant_weights(quant_method: QuantizeMethodBase | None) -> bool:
    """Returns whether the quantization method should load quantized weights."""
    return quant_method is not None and not isinstance(
        quant_method, UnquantizedLinearMethod
    )


def set_default_quant_scales(layer: nn.Module, register_buffer: bool = False) -> None:
    """Sets default quantization scales for the layer."""
    if register_buffer:
        layer.register_buffer("_k_scale", torch.tensor(1.0, dtype=torch.float32))
        layer.register_buffer("_v_scale", torch.tensor(1.0, dtype=torch.float32))
        layer.register_buffer("_q_scale", torch.tensor(1.0, dtype=torch.float32))
        layer.register_buffer("_prob_scale", torch.tensor(1.0, dtype=torch.float32))
    else:
        layer._k_scale.fill_(1.0)
        layer._v_scale.fill_(1.0)
        layer._q_scale.fill_(1.0)
        layer._prob_scale.fill_(1.0)

    # 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
    layer._prob_scale_float = 1.0

115
116
117
118
119
    # Initialize q/k/v range constants used by calc_kv_scales
    layer.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
    layer.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
    layer.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

120

121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
def _init_kv_cache_quant(
    layer: nn.Module,
    quant_config: QuantizationConfig | None,
    prefix: str,
) -> 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.
    """

139
140
141
142
143
144
145
146
147
148
149
150
151
152
    # Note [Register q/k/v/prob scales in state dict]
    # When calling model.to(device), only parameters/buffers in state dict are
    # moved. If not registering q/k/v/prob scales in state dict, there would
    # be an IMA error when a cuda kernel (e.g., quant_fp8) accesses the tensor
    # on cpu.
    # Registering in state dict means it interacts with weight loading. One edge
    # case is when quant_method is None, or quant_method is UnquantizedLinearMethod
    # (i.e., should_load_quant_weights(quant_method) == False).
    # In this case, the checkpoint does not have the scales. We need to
    # initialize the scales to 1.0 and update the scales after weight loading.
    # This is espectially important when we load dummy weights first (providing
    # wrong scales) and then load real weights (which misses scales and keeps the
    # wrong scales from dummy load).
    set_default_quant_scales(layer, register_buffer=True)
153
154
155
156
157
158
159
160

    # 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
    )
161
162
163

    # See [Note: Register q/k/v/prob scales in state dict]
    if should_load_quant_weights(quant_method):
164
165
166
        assert isinstance(quant_method, BaseKVCacheMethod)
        # TODO (mgoin): kv cache dtype should be specified in the FP8
        # checkpoint config and become the "auto" behavior
167
        if layer.kv_cache_dtype == "fp8_e5m2":
168
169
170
171
172
173
174
175
176
            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)


177
class Attention(nn.Module, AttentionLayerBase):
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
    """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,
194
195
        num_kv_heads: int | None = None,
        alibi_slopes: list[float] | None = None,
Li Xie's avatar
Li Xie committed
196
        use_alibi_sqrt: bool | None = None,
197
198
199
200
        cache_config: CacheConfig | None = None,
        quant_config: QuantizationConfig | None = None,
        logits_soft_cap: float | None = None,
        per_layer_sliding_window: int | None = None,
201
        prefix: str = "",
202
        attn_type: str = AttentionType.DECODER,
203
204
        kv_sharing_target_layer_name: str | None = None,
        attn_backend: type[AttentionBackend] | None = None,
205
        head_size_v: int | None = None,
206
        **extra_impl_args,
207
    ) -> None:
208
209
210
211
        """
        The KV cache is stored inside this class and is accessed via
        `self.kv_cache`.
        """
212
        super().__init__()
213
        sliding_window: int | None
214
215
216
217
218
219
220
221
222
        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

223
        vllm_config = get_current_vllm_config()
224
225
        if cache_config is not None:
            kv_cache_dtype = cache_config.cache_dtype
226
            calculate_kv_scales = cache_config.calculate_kv_scales
227
228
        else:
            kv_cache_dtype = "auto"
229
            calculate_kv_scales = False
230
231

        # llm-compressor mdls need to set cache_dtype to "fp8" manually.
232
233
        kv_cache_scheme = getattr(quant_config, "kv_cache_scheme", None)
        if kv_cache_scheme is not None:
234
235
236
237
238
239
            kv_cache_dtype = "fp8"
            calculate_kv_scales = False
            if cache_config is not None:
                cache_config.cache_dtype = "fp8"
                cache_config.calculate_kv_scales = False

240
241
242
243
244
245
        # Check if per-head quant scales are required based on kv_cache_scheme
        use_per_head_quant_scales = (
            kv_cache_scheme is not None
            and kv_cache_scheme.get("strategy") == "attn_head"
        )

246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        # Skip quantization for specified layers
        if cache_config is not None and cache_config.kv_cache_dtype_skip_layers:
            from vllm.model_executor.models.utils import extract_layer_index

            skip = False
            # Check attention type
            if (
                sliding_window is not None
                and "sliding_window" in cache_config.kv_cache_dtype_skip_layers
            ):
                skip = True
            # Check layer index
            layer_idx = extract_layer_index(prefix)
            if str(layer_idx) in cache_config.kv_cache_dtype_skip_layers:
                skip = True
            if skip:
                kv_cache_dtype = "auto"
                calculate_kv_scales = False
264
            logger.debug(
265
266
267
268
269
270
                "Layer %s: kv_cache_dtype=%s, sliding_window=%s",
                prefix,
                kv_cache_dtype,
                sliding_window,
            )

271
272
273
        self.kv_cache_torch_dtype = kv_cache_dtype_str_to_dtype(
            kv_cache_dtype, vllm_config.model_config
        )
274
275
        self.kv_cache_dtype = kv_cache_dtype
        self.calculate_kv_scales = calculate_kv_scales
276
277
        if num_kv_heads is None:
            num_kv_heads = num_heads
278
279
280
        assert num_heads % num_kv_heads == 0, (
            f"num_heads ({num_heads}) is not divisible by num_kv_heads ({num_kv_heads})"
        )
281
282
        self.quant_config = quant_config
        self.layer_name = prefix
283

284
285
        self.num_heads = num_heads
        self.head_size = head_size
286
        self.head_size_v = self.head_size if head_size_v is None else head_size_v
287
288
        self.num_kv_heads = num_kv_heads
        self.sliding_window = sliding_window
289
        self.has_sink = extra_impl_args.get("sinks") is not None
290

291
292
293
294
        # NOTE: model_config may be None during certain tests
        model_config = vllm_config.model_config
        self.use_mm_prefix = model_config is not None and model_config.is_mm_prefix_lm

295
296
297
        # During model initialization, the default dtype is set as the model
        # weight and activation dtype.
        dtype = torch.get_default_dtype()
298
        if attn_backend is None:
299
300
301
302
            self.attn_backend = get_attn_backend(
                head_size,
                dtype,
                kv_cache_dtype,
303
                use_mla=False,
304
                has_sink=self.has_sink,
305
                use_mm_prefix=self.use_mm_prefix,
306
                use_per_head_quant_scales=use_per_head_quant_scales,
307
                attn_type=attn_type,
308
            )
309
310
        else:
            self.attn_backend = attn_backend
Li Xie's avatar
Li Xie committed
311
312
313
314
315
316
317
318
319
320
        backend_supports_alibi_sqrt = self.attn_backend.supports_alibi_sqrt()
        use_alibi_sqrt = use_alibi_sqrt if use_alibi_sqrt else False
        if use_alibi_sqrt and not backend_supports_alibi_sqrt:
            raise ValueError(
                f"use_alibi_sqrt is not supported by backend "
                f"{self.attn_backend.get_name()}."
            )
        self.use_alibi_sqrt = bool(use_alibi_sqrt)
        if backend_supports_alibi_sqrt:
            extra_impl_args["use_alibi_sqrt"] = self.use_alibi_sqrt
321
322
323
324
325
        # prefix caching + batch invariance is currently not supported for
        # FLASHINFER and TRITON_MLA.
        if (
            cache_config is not None
            and cache_config.enable_prefix_caching
326
            and envs.VLLM_BATCH_INVARIANT
327
328
329
330
331
332
333
334
335
336
337
            and (
                self.attn_backend.get_name() == "FLASHINFER"
                or self.attn_backend.get_name() == "TRITON_MLA"
            )
        ):
            logger.warning_once(
                "Disabling prefix caching for FLASHINFER/TRITON_MLA "
                "with batch invariance, as it is not yet supported.",
            )
            cache_config.enable_prefix_caching = False

338
339
340
341
342
343
        if extra_impl_args.get("chunk_lookback", -1) > -1:
            assert self.attn_backend.get_name() == "TRITON_ATTN", (
                f"Chunked attention with lookback requires the Triton backend, "
                f"but got {self.attn_backend.get_name()}."
            )

344
        impl_cls = self.attn_backend.get_impl_cls()
345
        self.impl = impl_cls(  # type: ignore[assignment]  # impl_cls always returns an AttentionImpl subclass
346
347
348
349
350
351
352
353
354
355
356
357
            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,
        )
358
        self.backend = AttentionBackendEnum[self.attn_backend.get_name()]
359
        self.dtype = dtype
360

361
362
363
364
        # 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.
365
        self.use_direct_call = not current_platform.opaque_attention_op()
366

367
        compilation_config = vllm_config.compilation_config
368
369
370
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self
371
        self.attn_type = attn_type
372
373
374
375
376
377
378
379
380

        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

381
382
383
        # 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
384
        self.kv_cache = torch.tensor([])
385

386
387
        # Initialize KV cache quantization attributes
        _init_kv_cache_quant(self, quant_config, prefix)
388

389
390
391
392
        # Initialize TurboQuant buffers (Pi, S, centroids) if tq cache dtype
        if kv_cache_dtype.startswith("turboquant_"):
            self._init_turboquant_buffers(kv_cache_dtype, head_size, prefix)

393
394
        # for attn backends supporting query quantization
        self.query_quant = None
395
396
        if (
            self.impl.supports_quant_query_input
397
398
399
            and (
                self.kv_cache_dtype.startswith("fp8") or self.kv_cache_dtype == "nvfp4"
            )
400
            and not self.kv_cache_dtype.endswith("per_token_head")
401
        ):
402
403
404
405
406
407
408
409
410
411
            is_per_head = (
                hasattr(self, "q_scale") and self.q_scale.numel() == self.num_kv_heads
            )
            block_size = self.head_size * self.num_heads // self.num_kv_heads
            self.query_quant = QuantFP8(
                static=True,
                group_shape=GroupShape(-1, block_size)
                if is_per_head
                else GroupShape.PER_TENSOR,
            )
412

413
414
415
    def _init_turboquant_buffers(
        self, cache_dtype: str, head_size: int, prefix: str
    ) -> None:
416
        """Initialize TurboQuant centroids for Lloyd-Max quantization."""
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
        from vllm.model_executor.layers.quantization.turboquant.centroids import (
            get_centroids,
        )
        from vllm.model_executor.layers.quantization.turboquant.config import (
            TurboQuantConfig,
        )

        tq_config = TurboQuantConfig.from_cache_dtype(cache_dtype, head_size)

        self.register_buffer(
            "_tq_centroids",
            get_centroids(head_size, tq_config.centroid_bits),
        )
        self._tq_config = tq_config

        # Pre-allocate decode intermediate buffers so model.to(device) moves
        # them to GPU *before* the memory profiler runs.  Without this the
        # profiler gives all free memory to KV cache blocks and the first
        # decode OOMs when these buffers are lazily allocated.
        _vllm_cfg = get_current_vllm_config()
        B = _vllm_cfg.scheduler_config.max_num_seqs
        Hq = self.num_heads
        S = _vllm_cfg.attention_config.tq_max_kv_splits_for_cuda_graph
        D = head_size
        self.register_buffer(
            "_tq_mid_o_buf",
            torch.empty(B, Hq, S, D + 1, dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_tq_output_buf",
            torch.empty(B, Hq, D, dtype=torch.float32),
            persistent=False,
        )
        self.register_buffer(
            "_tq_lse_buf",
            torch.empty(B, Hq, dtype=torch.float32),
            persistent=False,
        )

457
458
459
460
461
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
462
463
464
        # 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.
465
        output_shape: torch.Size | None = None,
466
    ) -> torch.Tensor:
467
468
469
470
471
472
473
474
475
        """
        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
476
        if self.calculate_kv_scales:
477
478
479
            torch.ops.vllm.maybe_calc_kv_scales(
                query, key, value, _encode_layer_name(self.layer_name)
            )
480
481
482
483
484
485
486
        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
487
            assert self.kv_cache_dtype in {"fp8", "fp8_e4m3", "nvfp4"}
488
489

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

493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
        if output_shape is None:
            # Handle both 2D [num_tokens, hidden] and
            # 3D [num_tokens, heads, head_dim] query
            num_tokens = query.shape[0]
            output_shape = torch.Size((num_tokens, self.num_heads * self.head_size_v))
        output = torch.empty(output_shape, dtype=output_dtype, device=query.device)
        hidden_size = output_shape[-1]
        # 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_v)
        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_v)
        kv_cache_dummy_dep = None
        if self.use_direct_call:
            # Skip this if sharing KV cache with an earlier attention layer.
            if (
                not self.attn_backend.forward_includes_kv_cache_update
                and self.kv_sharing_target_layer_name is None
                and key is not None
                and value is not None
            ):
                kv_cache_dummy_dep = unified_kv_cache_update(
                    key, value, self.layer_name
520
                )
521
522
523
524
525
526
527
            unified_attention_with_output(
                query,
                key,
                value,
                output,
                self.layer_name,
                kv_cache_dummy_dep=kv_cache_dummy_dep,
528
            )
529
530
        else:
            # Skip this if sharing KV cache with an earlier attention layer.
531
            encoded = _encode_layer_name(self.layer_name)
532
533
534
535
536
537
538
            if (
                not self.attn_backend.forward_includes_kv_cache_update
                and self.kv_sharing_target_layer_name is None
                and key is not None
                and value is not None
            ):
                kv_cache_dummy_dep = torch.ops.vllm.unified_kv_cache_update(
539
                    key, value, encoded
540
                )
541
542
543
544
545
            torch.ops.vllm.unified_attention_with_output(
                query,
                key,
                value,
                output,
546
                encoded,
547
548
549
                kv_cache_dummy_dep=kv_cache_dummy_dep,
            )
        return output.view(-1, hidden_size)
550

551
552
    def calc_kv_scales(self, query, key, value):
        self._q_scale.copy_(torch.abs(query).max() / self.q_range)
553
554
        self._k_scale.copy_(torch.abs(key).max() / self.k_range)
        self._v_scale.copy_(torch.abs(value).max() / self.v_range)
555
        self._q_scale_float = self._q_scale.item()
556
557
558
559
560
        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

561
562
563
564
565
    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
566
        s += f", backend={self.impl.__class__.__name__}"
567
        return s
568

569
    def process_weights_after_loading(self, act_dtype: torch.dtype):
570
        self.impl.process_weights_after_loading(act_dtype)
571

572
573
574
575
576
577
578
579
580
581
582
        # If we should not load quant weights, we initialize the scales to 1.0
        # as the default value. See [Note: Register q/k/v/prob scales in state dict]
        # for more details.
        quant_method = (
            self.quant_config.get_quant_method(self, prefix=self.layer_name)
            if self.quant_config
            else None
        )
        if not should_load_quant_weights(quant_method):
            set_default_quant_scales(self, register_buffer=False)

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

586
    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
587
588
589
590
        # 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
591
        quant_mode = get_kv_quant_mode(self.kv_cache_dtype)
592
593
594
595
596
597
598
599
        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,
600
                head_size_v=self.head_size_v,
601
                dtype=self.kv_cache_torch_dtype,
602
                kv_quant_mode=quant_mode,
603
604
                sliding_window=self.sliding_window,
            )
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
        elif self.kv_cache_dtype.startswith("turboquant_"):
            from vllm.model_executor.layers.quantization.turboquant.config import (
                TurboQuantConfig,
            )
            from vllm.v1.kv_cache_interface import TQFullAttentionSpec

            tq_config = TurboQuantConfig.from_cache_dtype(
                self.kv_cache_dtype, self.head_size
            )
            return TQFullAttentionSpec(
                block_size=block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
                head_size_v=self.head_size,
                dtype=self.kv_cache_torch_dtype,
                tq_slot_size=tq_config.slot_size_aligned,
            )
622
623
624
625
626
        else:
            return FullAttentionSpec(
                block_size=block_size,
                num_kv_heads=self.num_kv_heads,
                head_size=self.head_size,
627
                head_size_v=self.head_size_v,
628
                dtype=self.kv_cache_torch_dtype,
629
                kv_quant_mode=quant_mode,
630
631
            )

632

633
634
635
636
def maybe_calc_kv_scales(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
637
    layer_name: LayerNameType,
638
) -> None:
639
    layer_name = _resolve_layer_name(layer_name)
640
    forward_context: ForwardContext = get_forward_context()
641
    self = forward_context.no_compile_layers[layer_name]
642

643
644
645
    # 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:
646
647
648
649
650
651
652
653
654
        return

    self.calc_kv_scales(query, key, value)


def maybe_calc_kv_scales_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
655
    layer_name: LayerNameType,
656
657
658
659
660
661
662
663
664
665
666
667
) -> 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,
)


668
def get_attention_context(
669
    layer_name: str,
670
) -> tuple[Any, "Attention | MLAAttention", torch.Tensor, torch.Tensor]:
671
672
673
    """Extract attention context for a given layer.

    This helper function extracts the attention metadata, attention layer
674
    instance, KV cache tensor, and slot mapping for a specific layer.
675
676
677

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

679
680
681
682
683
    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)
684
        - kv_cache: The KV cache tensor for current forward pass
685
        - slot_mapping: The slot mapping for this specific layer
686
687
688
689

        Note: attn_metadata may be None, but attn_layer and kv_cache are always
        extracted from the forward context.
    """
690
    forward_context: ForwardContext = get_forward_context()
691
692
693
694
695
696
697
698
699
700
    attn_metadata_raw = forward_context.attn_metadata
    attn_metadata: AttentionMetadata
    if isinstance(attn_metadata_raw, dict):
        attn_metadata = attn_metadata_raw[layer_name]
    elif isinstance(attn_metadata_raw, list):
        # list[dict[str, AttentionMetadata]]: used in speculative decoding
        # where [0] is the base-model (non-speculative) metadata dict.
        attn_metadata = attn_metadata_raw[0][layer_name]
    else:
        attn_metadata = attn_metadata_raw
701
    attn_layer: Attention | MLAAttention = forward_context.no_compile_layers[layer_name]
702
    kv_cache = attn_layer.kv_cache
703
704
705
706
707
708
    slot_mapping = forward_context.slot_mapping
    assert isinstance(slot_mapping, dict), (
        f"Expected slot_mapping to be a dict, got {type(slot_mapping)}. "
    )
    layer_slot_mapping = slot_mapping.get(layer_name)
    return attn_metadata, attn_layer, kv_cache, layer_slot_mapping
709
710


711
712
713
def unified_kv_cache_update(
    key: torch.Tensor,
    value: torch.Tensor,
714
    layer_name: LayerNameType,
715
716
717
718
719
) -> torch.Tensor:
    """
    Returns a dummy that is passed to unified_attention to signal a side effect and
    the data dependency between them to ensure torch.compile preserves ordering.
    """
720
    layer_name = _resolve_layer_name(layer_name)
721
    _, attn_layer, kv_cache, layer_slot_mapping = get_attention_context(layer_name)
722
723
724
725
    if layer_slot_mapping is not None:
        assert hasattr(attn_layer.impl, "do_kv_cache_update"), (
            f"{attn_layer.impl.__class__.__name__} does not support kv cache update"
        )
726
        attn_layer.impl.do_kv_cache_update(  # type: ignore[attr-defined]
727
728
729
730
731
732
733
734
735
736
737
738
739
            attn_layer,
            key,
            value,
            kv_cache,
            layer_slot_mapping,
        )

    return torch.empty(0, device=kv_cache.device, dtype=kv_cache.dtype)


def unified_kv_cache_update_fake(
    key: torch.Tensor,
    value: torch.Tensor,
740
    layer_name: LayerNameType,
741
742
743
744
745
746
747
748
749
750
751
752
) -> torch.Tensor:
    return torch.empty(0, device=key.device, dtype=key.dtype)


direct_register_custom_op(
    op_name="unified_kv_cache_update",
    op_func=unified_kv_cache_update,
    fake_impl=unified_kv_cache_update_fake,
    mutates_args=[],
)


753
@maybe_transfer_kv_layer
754
755
756
757
758
def unified_attention_with_output(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
759
    layer_name: LayerNameType,
760
761
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
762
    kv_cache_dummy_dep: torch.Tensor | None = None,
763
) -> None:
764
765
766
767
    # kv_cache_dummy_dep is not used but accepting it creates a data dependency
    # that ensures torch.compile preserves ordering between KV cache update and
    # attention forward.
    del kv_cache_dummy_dep
768
    layer_name = _resolve_layer_name(layer_name)
769
    attn_metadata, self, kv_cache, _ = get_attention_context(layer_name)
770

771
772
773
774
775
776
777
778
779
780
781
    self.impl.forward(
        self,
        query,
        key,
        value,
        kv_cache,
        attn_metadata,
        output=output,
        output_scale=output_scale,
        output_block_scale=output_block_scale,
    )
782
783
784
785
786
787
788


def unified_attention_with_output_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
789
    layer_name: LayerNameType,
790
791
    output_scale: torch.Tensor | None = None,
    output_block_scale: torch.Tensor | None = None,
792
    kv_cache_dummy_dep: torch.Tensor | None = None,
793
794
795
796
797
798
799
) -> None:
    return


direct_register_custom_op(
    op_name="unified_attention_with_output",
    op_func=unified_attention_with_output,
800
    mutates_args=["output", "output_block_scale"],
801
802
    fake_impl=unified_attention_with_output_fake,
)