layer.py 23.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Attention layer."""
4
from typing import Any, Dict, List, Optional
5
6
7

import torch
import torch.nn as nn
8
import torch.nn.functional as F
9

10
from vllm.two_batch_overlap.v1.two_batch_overlap_v1 import tbo_maybe_save_kv_layer_to_connector
11
import vllm.envs as envs
12
from vllm.attention import AttentionType
13
from vllm.attention.selector import backend_name_to_enum, get_attn_backend
14
from vllm.config import CacheConfig, get_current_vllm_config
15
16
17
from vllm.distributed.kv_transfer import (get_kv_transfer_group,
                                          has_kv_transfer_group,
                                          is_v1_kv_transfer_group)
18
from vllm.forward_context import ForwardContext, get_forward_context
19
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
20
21
from vllm.model_executor.layers.quantization.base_config import (
    QuantizationConfig)
22
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
23
from vllm.platforms import _Backend, current_platform
24
from vllm.utils import direct_register_custom_op
25
from vllm.v1.attention.backends.utils import validate_kv_sharing_target
26

zhuwenwen's avatar
zhuwenwen committed
27
28
29
30
31
USE_XFORMERS_OPS = None
try:
    tag_cudagraph_unsafe = (torch._C.Tag.cudagraph_unsafe, )
except AttributeError:
    tag_cudagraph_unsafe = ()  # type: ignore[assignment]
32
    
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

class Attention(nn.Module):
    """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,
        num_kv_heads: Optional[int] = None,
        alibi_slopes: Optional[List[float]] = None,
53
        cache_config: Optional[CacheConfig] = None,
54
        quant_config: Optional[QuantizationConfig] = None,
55
        blocksparse_params: Optional[Dict[str, Any]] = None,
56
        logits_soft_cap: Optional[float] = None,
57
        per_layer_sliding_window: Optional[int] = None,
58
        use_mla: bool = False,
59
        prefix: str = "",
60
        attn_type: str = AttentionType.DECODER,
61
        kv_sharing_target_layer_name: Optional[str] = None,
62
        **extra_impl_args,
63
    ) -> None:
64
65
66
67
        """
        The KV cache is stored inside this class and is accessed via
        `self.kv_cache`.
        """
68
        super().__init__()
69
70
71
72
73
74
75
76
77
        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

78
79
80
        if cache_config is not None:
            kv_cache_dtype = cache_config.cache_dtype
            block_size = cache_config.block_size
81
            is_attention_free = cache_config.is_attention_free
82
            calculate_kv_scales = cache_config.calculate_kv_scales
83
84
        else:
            kv_cache_dtype = "auto"
85
            block_size = 64 if envs.VLLM_USE_FLASH_ATTN_PA and envs.VLLM_USE_FLASH_MLA else 16
86
            is_attention_free = False
87
            calculate_kv_scales = False
88
89
        if num_kv_heads is None:
            num_kv_heads = num_heads
90
91
92
        assert num_heads % num_kv_heads == 0, \
            f"num_heads ({num_heads}) is not " \
            f"divisible by num_kv_heads ({num_kv_heads})"
93

94
        # The default k/v_scale is set to 1.0. This is ignored
95
96
        # when kv-cache is not fp8, and should be used with
        # kv-cache in fp8_e5m2. For kv-cache in fp8_e4m3, we
97
        # expect the pre-quantized k/v_scale to be loaded along
98
99
        # with the model weights.
        self.kv_cache_dtype = kv_cache_dtype
100
        self.calculate_kv_scales = calculate_kv_scales
101
102
103
        if torch.cuda.get_device_properties("cuda").gcnArchName.split(':')[0] == "gfx938" and kv_cache_dtype == "fp8_e4m3" and envs.VLLM_USE_FLASH_MLA_FP8:
            self._k_scale = torch.ones((1), dtype=torch.float32)
            self._v_scale = torch.ones((1), dtype=torch.float32)
zhuwenwen's avatar
zhuwenwen committed
104
            self._q_scale = torch.ones((1), dtype=torch.float32)
105
106
107
        else:
            self._k_scale = torch.tensor(1.0, dtype=torch.float32)
            self._v_scale = torch.tensor(1.0, dtype=torch.float32)
zhuwenwen's avatar
zhuwenwen committed
108
            self._q_scale = torch.tensor(1.0, dtype=torch.float32)
109
110
        # FlashAttn doesn't support quantizing the kv-cache only
        # but requires q to be quantized as well.
111
        self._prob_scale = torch.tensor(1.0, dtype=torch.float32)
112
113
114
115
116
117

        # We also keep the float32 versions of k/v_scale for attention
        # backends that don't support tensors (Flashinfer)
        self._k_scale_float = 1.0
        self._v_scale_float = 1.0

118
        self.use_mla = use_mla
119
120
121
122
123
        self.num_heads = num_heads
        self.head_size = head_size
        self.num_kv_heads = num_kv_heads
        self.sliding_window = sliding_window

124
        quant_method = quant_config.get_quant_method(
125
            self, prefix=prefix) if quant_config else None
126
127
        if quant_method is not None and not isinstance(
                quant_method, UnquantizedLinearMethod):
128
            assert isinstance(quant_method, BaseKVCacheMethod)
129
130
            # TODO (mgoin): kv cache dtype should be specified in the FP8
            # checkpoint config and become the "auto" behavior
zhuwenwen's avatar
zhuwenwen committed
131
132
133
            # if self.kv_cache_dtype == "fp8_e5m2":
            #     raise ValueError("fp8_e5m2 kv-cache is not supported with "
            #                      "fp8 checkpoints.")
134
135
136
137
138
139
            # 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.
            self.quant_method = quant_method
            self.quant_method.create_weights(self)
140

141
142
143
        # During model initialization, the default dtype is set as the model
        # weight and activation dtype.
        dtype = torch.get_default_dtype()
144
145
146
147
148
149
150
        attn_backend = get_attn_backend(head_size,
                                        dtype,
                                        kv_cache_dtype,
                                        block_size,
                                        is_attention_free,
                                        blocksparse_params is not None,
                                        use_mla=use_mla)
151
        impl_cls = attn_backend.get_impl_cls()
152
        self.impl = impl_cls(num_heads, head_size, scale, num_kv_heads,
153
                             alibi_slopes, sliding_window, kv_cache_dtype,
154
                             blocksparse_params, logits_soft_cap, attn_type,
155
                             kv_sharing_target_layer_name, **extra_impl_args)
156
        self.backend = backend_name_to_enum(attn_backend.get_name())
157
        self.dtype = dtype
158

159
160
161
162
        # 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.
163
164
165
        self.use_direct_call = not current_platform.is_cuda_alike(
        ) and not current_platform.is_cpu()

166
        self.use_output = attn_backend.accept_output_buffer
167
168
169
170
171
        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.layer_name = prefix
172
        self.attn_type = attn_type
173
174
175
176
177
178
179
180
181
182
183
184
185

        if kv_sharing_target_layer_name is not None:
            if not envs.VLLM_USE_V1:
                raise NotImplementedError(
                    "Cross-layer KV sharing is not supported in V0.")

            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

186
187
188
189
190
191
192
        # 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 = [
            torch.tensor([]) for _ in range(get_current_vllm_config(
            ).parallel_config.pipeline_parallel_size)
        ]
193

194
        self.q_range = torch.tensor(envs.Q_SCALE_CONSTANT, dtype=torch.float32)
195
196
197
        self.k_range = torch.tensor(envs.K_SCALE_CONSTANT, dtype=torch.float32)
        self.v_range = torch.tensor(envs.V_SCALE_CONSTANT, dtype=torch.float32)

198
199
200
201
202
    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
203
204
205
206
        # 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.
        output_shape: Optional[torch.Size] = None,
207
208
        query_nope: Optional[torch.Size] = None,
        num_local_heads: Optional[int] = None,
209
210
211
212
213
        q_ori: Optional[torch.Tensor] = None,
        key_normed: Optional[torch.Tensor] = None,
        positions: Optional[torch.Tensor] = None,
        weight: Optional[torch.Tensor] = None,
        cos_sin_cache: Optional[torch.Tensor] = None,
214
    ) -> torch.Tensor:
215
216
217
218
219
220
221
222
223
        """
        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
224
        if self.calculate_kv_scales:
zhuwenwen's avatar
zhuwenwen committed
225
226
            # attn_metadata = get_forward_context().attn_metadata
            # #if (attn_metadata is not None and getattr(attn_metadata, "enable_kv_scales_calculation", False)):
yangql's avatar
yangql committed
227
            # if key is not None and value is not None:         
zhuwenwen's avatar
zhuwenwen committed
228
229
230
                # self.calc_kv_scales(query, key, value)
            torch.ops.vllm.maybe_calc_kv_scales(query, key, value,
                                                self.layer_name)     
231
        if self.use_output:
232
233
            output_shape = (output_shape
                            if output_shape is not None else query.shape)
234
235
236
237
238
239
240
241
            if envs.VLLM_USE_OPT_ZEROS:
                output = torch.empty(output_shape,
                                    dtype=query.dtype,
                                    device=query.device)
            else:
                output = torch.zeros(output_shape,
                                    dtype=query.dtype,
                                    device=query.device)
242
243
244
245
246
247
248
249
250
251
252
253
254
255
            hidden_size = output_shape[-1]
            # We skip reshaping query, key and value tensors for the MLA
            # backend since these tensors have different semantics and are
            # processed differently.
            if not self.use_mla:
                # 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)
256
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
257
                forward_context: ForwardContext = get_forward_context()
258
                attn_metadata = forward_context.attn_metadata
259
260
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
261
262
263
264
265
266
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
                self.impl.forward(self,
                                  query,
                                  key,
                                  value,
                                  self_kv_cache,
267
                                  attn_metadata,
Chen Zhang's avatar
Chen Zhang committed
268
                                  output=output)
269
            else:
270
271
272
273
274
                if not envs.VLLM_USE_LIGHTOP_RMS_ROPE_CONCAT:
                    torch.ops.vllm.unified_attention_with_output(
                        query, key, value, output, self.layer_name)
                else:
                    torch.ops.vllm.unified_attention_with_output(
275
                        query, key, value, output, self.layer_name, None, query_nope, num_local_heads, q_ori, key_normed, positions, weight, cos_sin_cache)
276
            return output.view(-1, hidden_size)
277
        else:
278
            if self.use_direct_call:
Chen Zhang's avatar
Chen Zhang committed
279
                forward_context = get_forward_context()
280
                attn_metadata = forward_context.attn_metadata
281
282
                if isinstance(attn_metadata, dict):
                    attn_metadata = attn_metadata[self.layer_name]
Chen Zhang's avatar
Chen Zhang committed
283
284
                self_kv_cache = self.kv_cache[forward_context.virtual_engine]
                return self.impl.forward(self, query, key, value,
285
                                         self_kv_cache, attn_metadata)
286
287
288
            else:
                return torch.ops.vllm.unified_attention(
                    query, key, value, self.layer_name)
289

290
291
    def calc_kv_scales(self, query, key, value):
        self._q_scale.copy_(torch.abs(query).max() / self.q_range)
292
293
294
295
296
297
298
        self._k_scale.copy_(torch.abs(key).max() / self.k_range)
        self._v_scale.copy_(torch.abs(value).max() / self.v_range)
        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

299
300
301
302
303
    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
304
        s += f", backend={self.impl.__class__.__name__}"
305
        return s
306

307
    def process_weights_after_loading(self, act_dtype: torch.dtype):
308
        if hasattr(self.impl, "process_weights_after_loading"):
309
            self.impl.process_weights_after_loading(act_dtype)
310

311

312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
class MultiHeadAttention(nn.Module):
    """Multi-headed attention without any cache, used for ViT."""

    def __init__(
        self,
        num_heads: int,
        head_size: int,
        scale: float,
        num_kv_heads: Optional[int] = None,
    ):
        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

328
329
330
        assert self.num_heads % self.num_kv_heads == 0, \
            f"num_heads ({self.num_heads}) is not " \
            f"divisible by num_kv_heads ({self.num_kv_heads})"
331
332
        self.num_queries_per_kv = self.num_heads // self.num_kv_heads

333
334
335
336
        dtype = torch.get_default_dtype()
        attn_backend = get_attn_backend(head_size,
                                        dtype,
                                        kv_cache_dtype=None,
337
                                        block_size=64 if envs.VLLM_USE_FLASH_ATTN_PA and envs.VLLM_USE_FLASH_MLA else 16,
338
                                        is_attention_free=False)
339
        backend = backend_name_to_enum(attn_backend.get_name())
340
341
342
343
        if current_platform.is_rocm():
            # currently, only torch_sdpa is supported on rocm
            self.attn_backend = _Backend.TORCH_SDPA
        else:
344
345
            if backend in (_Backend.FLASH_ATTN, _Backend.FLASH_ATTN_VLLM_V1,
                           _Backend.FLEX_ATTENTION):
346
                backend = _Backend.XFORMERS
347

348
349
350
            self.attn_backend = backend if backend in {
                _Backend.TORCH_SDPA, _Backend.XFORMERS, _Backend.PALLAS_VLLM_V1
            } else _Backend.TORCH_SDPA
351
352
353
354
355
356
357
358

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
        value: torch.Tensor,
    ) -> torch.Tensor:
        """Input shape: batch_size x seq_len x hidden_size"""
359
        # TODO(Isotr0py): Use existing backend implementations and support FA3
360
361
362
363
364
365
366
        bsz, q_len, _ = query.size()
        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)

367
368
369
370
371
        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)

372
        if self.attn_backend == _Backend.XFORMERS:
373
374
375
376
377
378
379
380
381
382
383
384
385
386
            from xformers import ops as xops

            out = xops.memory_efficient_attention_forward(query,
                                                          key,
                                                          value,
                                                          scale=self.scale)
        elif self.attn_backend == _Backend.TORCH_SDPA:
            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)
            out = out.transpose(1, 2)
387
388
389
390
391
392
393
        elif self.attn_backend == _Backend.PALLAS_VLLM_V1:
            query, key, value = (x.transpose(1, 2)
                                 for x in (query, key, value))
            from torch_xla.experimental.custom_kernel import flash_attention
            out = flash_attention(query, key, value, sm_scale=self.scale)
            out = out.transpose(1, 2)

394
        return out.reshape(bsz, q_len, -1)
395
396


397
398
399
400
401
402
403
404
405
406
def wait_for_kv_layer_from_connector(layer_name: str):
    if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
        return

    connector = get_kv_transfer_group()

    forward_context: ForwardContext = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    if attn_metadata is None:
        return
407
    assert isinstance(attn_metadata, dict)
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
    connector.wait_for_layer_load(layer_name)


def maybe_save_kv_layer_to_connector(
    layer_name: str,
    kv_cache_layer: List[torch.Tensor],
):
    if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
        return

    connector = get_kv_transfer_group()

    forward_context: ForwardContext = get_forward_context()
    attn_metadata = forward_context.attn_metadata
    if attn_metadata is None:
        return
424
425
426
    assert isinstance(attn_metadata, dict)
    connector.save_kv_layer(layer_name, kv_cache_layer,
                            attn_metadata[layer_name])
427

428

zhuwenwen's avatar
zhuwenwen committed
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
457
458
459
460
461
462
463
464
465
466
def maybe_calc_kv_scales(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> None:

    forward_context: ForwardContext = get_forward_context()
    attn_metadata = forward_context.attn_metadata

    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]

    # if attn_metadata is None or not getattr(
    #         attn_metadata, 'enable_kv_scales_calculation', False):
    #     return

    self = forward_context.no_compile_layers[layer_name]
    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=[],
    fake_impl=maybe_calc_kv_scales_fake,
    dispatch_key=current_platform.dispatch_key,
    tags=tag_cudagraph_unsafe,)


467
468
469
470
471
472
def unified_attention(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    layer_name: str,
) -> torch.Tensor:
473
474
    wait_for_kv_layer_from_connector(layer_name)

475
    forward_context: ForwardContext = get_forward_context()
476
    attn_metadata = forward_context.attn_metadata
477
478
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
479
    self = forward_context.no_compile_layers[layer_name]
480
    kv_cache = self.kv_cache[forward_context.virtual_engine]
481
482
483
    output = self.impl.forward(self, query, key, value, kv_cache,
                               attn_metadata)

484
485
486
487
    if envs.VLLM_ENABLE_TBO:
        tbo_maybe_save_kv_layer_to_connector(layer_name, kv_cache)
    else:
        maybe_save_kv_layer_to_connector(layer_name, kv_cache)
488
    return output
489
490
491
492
493
494
495
496
497
498
499
500
501
502


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,
503
    mutates_args=[],
504
505
506
    fake_impl=unified_attention_fake,
    dispatch_key=current_platform.dispatch_key,
)
507
508
509
510
511
512
513
514


def unified_attention_with_output(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
515
    output_scale: Optional[torch.Tensor] = None,
516
517
    query_nope: Optional[torch.Tensor] = None,
    num_local_heads: Optional[int] = None,
518
519
520
521
522
    q_ori: Optional[torch.Tensor] = None,
    key_normed: Optional[torch.Tensor] = None,
    positions: Optional[torch.Tensor] = None,
    weight: Optional[torch.Tensor] = None,
    cos_sin_cache: Optional[torch.Tensor] = None,
523
) -> None:
524
    wait_for_kv_layer_from_connector(layer_name)
525
    forward_context: ForwardContext = get_forward_context()
526
    attn_metadata = forward_context.attn_metadata
527
528
    if isinstance(attn_metadata, dict):
        attn_metadata = attn_metadata[layer_name]
529
    self = forward_context.no_compile_layers[layer_name]
530
    kv_cache = self.kv_cache[forward_context.virtual_engine]
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
    if not envs.VLLM_USE_LIGHTOP_RMS_ROPE_CONCAT:
        self.impl.forward(self,
                        query,
                        key,
                        value,
                        kv_cache,
                        attn_metadata,
                        output=output,
                        output_scale=output_scale)
    else:
        self.impl.forward(self,
                        query,
                        key,
                        value,
                        kv_cache,
                        attn_metadata,
                        output=output,
                        output_scale=output_scale,
549
550
                        query_nope=query_nope,
                        num_local_heads=num_local_heads,
551
552
553
554
555
                        q_ori=q_ori,
                        key_normed=key_normed,
                        positions=positions,
                        weight=weight,
                        cos_sin_cache=cos_sin_cache)
556
557
558
559
    if envs.VLLM_ENABLE_TBO:
        tbo_maybe_save_kv_layer_to_connector(layer_name, kv_cache)
    else:
        maybe_save_kv_layer_to_connector(layer_name, kv_cache)
560

561

562
563
564
565
566
567
568
569
570
571
572
573
def unified_attention_with_output_fake(
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    output: torch.Tensor,
    layer_name: str,
    output_scale: Optional[torch.Tensor] = None,
    q_ori: Optional[torch.Tensor] = None,
    key_normed: Optional[torch.Tensor] = None,
    positions: Optional[torch.Tensor] = None,
    weight: Optional[torch.Tensor] = None,
    cos_sin_cache: Optional[torch.Tensor] = None,
574
    ) -> None:
575
    return
576
577
578
579
580


direct_register_custom_op(
    op_name="unified_attention_with_output",
    op_func=unified_attention_with_output,
581
    mutates_args=["output"],
582
583
    fake_impl=unified_attention_with_output_fake,
    dispatch_key=current_platform.dispatch_key,
584
)