mslk-cuda-nightly 2026.1.19__cp310-cp310-manylinux_2_28_x86_64.whl

mslk/attention/fmha/common.py

@@ -0,0 +1,598 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# pyre-unsafe
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    List,
+    Mapping,
+    Optional,
+    Set,
+    Tuple,
+    Type,
+    Union,
+)
+
+import torch
+
+from .attn_bias import (
+    AttentionBias,
+    BlockDiagonalGappyKeysMask,
+    BlockDiagonalMask,
+    BlockDiagonalPaddedKeysMask,
+    LowerTriangularMask,
+    LowerTriangularMaskWithTensorBias,
+    PagedBlockDiagonalGappyKeysMask,
+    PagedBlockDiagonalPaddedKeysMask,
+)
+from .utils.cpp_lib import _built_with_cuda
+from .utils.op_common import BaseOperator
+
+
+def _is_bias_type_supported_in_BMK(attn_bias_type: Any) -> bool:
+    # NoneType
+    if isinstance(None, attn_bias_type):
+        return True
+    if attn_bias_type in [LowerTriangularMask, torch.Tensor]:
+        return True
+    return False
+
+
+def _attn_bias_apply(
+    attn_bias: Optional[Union[torch.Tensor, AttentionBias]],
+    op: Callable[[torch.Tensor], torch.Tensor],
+) -> Optional[Union[torch.Tensor, AttentionBias]]:
+    if isinstance(attn_bias, torch.Tensor):
+        return op(attn_bias)
+    if isinstance(attn_bias, LowerTriangularMaskWithTensorBias):
+        return LowerTriangularMaskWithTensorBias(op(attn_bias._bias))
+    return attn_bias
+
+
+class ScaledTensor(torch.Tensor):
+    __slots__ = ["scale", "dequant_func", "original_dtype"]
+
+    # Disabling custom torch function handling for this class
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    @staticmethod
+    def __new__(
+        cls,
+        data: torch.Tensor,
+        scale: torch.Tensor,
+        dequant_func: Callable[[torch.Tensor, torch.Tensor], torch.Tensor],
+        original_dtype: torch.dtype,
+        require_grad: bool = False,
+    ) -> "ScaledTensor":
+        """
+        Creates a new ScaledTensor subclass instance.
+
+        Parameters:
+        - data: The underlying quantized tensor (e.g., int8, int4).
+        - scale: The scale tensor or scalar to be used for dequantization.
+        - dequant_func: A callable that applies dequantization, which takes both the data and scale as input.
+        - original_dtype: The data type before quantization (e.g., float32, float16).
+        - require_grad: Whether or not to track gradients (default: False for inference use).
+        """
+        # Use _make_subclass to create a new ScaledTensor instance, which is a subclass of torch.Tensor.
+        instance = torch.Tensor._make_subclass(cls, data, require_grad)
+
+        # Store the dequantization scale and function as attributes.
+        instance.scale = scale  # type: ignore
+        instance.dequant_func = dequant_func  # type: ignore
+
+        # Store the original data type of the tensor, so we can cast it back after dequantization.
+        instance.original_dtype = original_dtype  # type: ignore
+
+        # Return the new instance of ScaledTensor.
+        return instance
+
+    def dequantize(self) -> torch.Tensor:
+        """
+        Applies the custom dequantization function provided at the tensor's creation.
+        After dequantization, the data is cast back to its original data type.
+        """
+        # Explicitly create a new torch.Tensor to ensure the return type is torch.Tensor, not ScaledTensor.
+        data = torch.Tensor(self.float())
+
+        # Call the dequantization function, passing in the data and the scale.
+        dequantized_data = self.dequant_func(data, self.scale)  # type: ignore
+
+        # Cast the dequantized data back to the original data type.
+        return dequantized_data.to(self.original_dtype)  # type: ignore
+
+    def unpack(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Unpacks the ScaledTensor by returning its data and scale as a tuple.
+        Returns:
+        - A tuple of (data, scale), both of which are torch.Tensor objects.
+        """
+        return self.data, self.scale  # type: ignore
+
+    def __repr__(self):
+        """
+        Custom string representation for ScaledTensor.
+        """
+        return f"ScaledTensor(data={self.data}, scale={self.scale}, original_dtype={self.original_dtype})"
+
+
+def pack_fp8_tensorwise_per_head(
+    x: torch.Tensor, scale: Union[torch.Tensor, float], original_dtype
+) -> ScaledTensor:
+    """
+    Pack a tensor into a tensorwise fp8 ScaledTensor.
+    """
+    if isinstance(scale, float):
+        scale = torch.tensor([scale], device=x.device)
+
+    def dequant_func(x, scale):
+        return x * scale[:, None, :, None]
+
+    return ScaledTensor(
+        data=x,
+        scale=scale,
+        dequant_func=dequant_func,
+        original_dtype=original_dtype,
+    )
+
+
+@dataclass
+class Inputs:
+    """
+    Stores inputs to the `memory_efficient_attention` operators
+    """
+
+    query: torch.Tensor
+    key: torch.Tensor
+    value: torch.Tensor
+    attn_bias: Optional[Union[torch.Tensor, AttentionBias]] = None
+    p: float = 0.0
+    scale: Optional[float] = None
+    output_dtype: Optional[torch.dtype] = None
+    is_partial: bool = False
+    quantize_pv_to_fp8: bool = False
+    quantize_qk_to_fp8: bool = False
+    use_fp32_scales: bool = False
+
+    @property
+    def device(self) -> torch.device:
+        return self.query.device
+
+    @property
+    def scale_float(self) -> float:
+        return self.query.shape[-1] ** (-0.5) if self.scale is None else self.scale
+
+    def get_qkv_in_bmghk(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        if self.query.ndim == 5:
+            return self.query, self.key, self.value
+        if self.query.ndim == 4:
+            return (
+                self.query.unsqueeze(2),
+                self.key.unsqueeze(2),
+                self.value.unsqueeze(2),
+            )
+        if self.value.ndim == 3:
+            return (
+                self.query[:, :, None, None],
+                self.key[:, :, None, None],
+                self.value[:, :, None, None],
+            )
+        raise AssertionError
+
+    def normalize_bmhk(self) -> Tuple[int, ...]:
+        if self.query.ndim not in [3, 4, 5]:
+            raise ValueError(
+                f"Invalid shape for query: {self.query.shape}. "
+                "Expected shape [batch, seqlen, head_groups, num_heads_per_group, K]"
+                ", [batch, seqlen, num_heads, K], or [batch, seqlen, K]."
+            )
+        if self.value.dtype == torch.int32:
+            # Quantized K/V case, in which the last dims of Q and K are different.
+            # NB we currently don't have any implementations for quantized KV with
+            # SUPPORTS_DIFFERENT_VALUE_EMBED.
+            output_shape: Tuple[int, ...] = tuple(self.query.shape)
+        else:
+            output_shape = tuple(self.query.shape)[:-1] + (self.value.shape[-1],)
+        # Convert from legacy format
+        if self.query.ndim == 3:
+            self.query = self.query.unsqueeze(2)
+            self.key = self.key.unsqueeze(2)
+            self.value = self.value.unsqueeze(2)
+            self.attn_bias = _attn_bias_apply(
+                self.attn_bias, partial(torch.unsqueeze, dim=1)
+            )
+        return output_shape
+
+    def validate_inputs(self) -> None:  # noqa: C901
+        qkv = (self.query, self.key, self.value)
+        if self.query.ndim not in (3, 4, 5) or any(
+            x.ndim != self.query.ndim for x in qkv
+        ):
+            raise ValueError(
+                f"Query/Key/Value should all have BMGHK, BMHK or BMK shape.\n"
+                f"  query.shape: {self.query.shape}\n"
+                f"  key.shape  : {self.key.shape}\n"
+                f"  value.shape: {self.value.shape}"
+            )
+        if any(x.device != self.query.device for x in qkv):
+            raise ValueError("Query/Key/Value should all be on the same device")
+        if isinstance(
+            self.attn_bias,
+            (
+                BlockDiagonalMask,
+                BlockDiagonalPaddedKeysMask,
+                PagedBlockDiagonalPaddedKeysMask,
+                BlockDiagonalGappyKeysMask,
+                PagedBlockDiagonalGappyKeysMask,
+            ),
+        ):
+            bias_device = self.attn_bias.q_seqinfo.seqstart.device
+            if bias_device != self.query.device:
+                raise ValueError(
+                    f"Attention bias and Query/Key/Value should be on the same device\n"
+                    f"  query.device: {self.query.device}\n"
+                    f"  attn_bias   : {bias_device}\n"
+                )
+
+        quantized_dtypes = self.key.dtype == self.value.dtype == torch.int32
+        non_quantized_dtypes = all(x.dtype == self.query.dtype for x in qkv)
+        if not (quantized_dtypes or non_quantized_dtypes):
+            raise ValueError(
+                "Query/Key/Value should either all have the same dtype, or "
+                "(in the quantized case) Key/Value should have dtype torch.int32\n"
+                f"  query.dtype: {self.query.dtype}\n"
+                f"  key.dtype  : {self.key.dtype}\n"
+                f"  value.dtype: {self.value.dtype}"
+            )
+        # Biases with tensors attached are meant to be in BMHK format
+        # This would require to permute biases/gradients which can be expensive,
+        # so let's just forbid it - BMK is a legacy format anyway
+        if self.query.ndim == 3 and not _is_bias_type_supported_in_BMK(
+            type(self.attn_bias)
+        ):
+            raise ValueError(
+                f"Please provide inputs in BMHK format rather "
+                f"than BMK when using bias type `{type(self.attn_bias).__name__}`"
+            )
+        attn_bias_t: Optional[torch.Tensor] = None
+        if isinstance(self.attn_bias, LowerTriangularMaskWithTensorBias):
+            attn_bias_t = self.attn_bias._bias
+        elif isinstance(self.attn_bias, torch.Tensor):
+            attn_bias_t = self.attn_bias
+        if self.query.ndim == 4 and attn_bias_t is not None:
+            expected_shape = (
+                self.query.shape[0],
+                self.query.shape[2],
+                self.query.shape[1],
+                self.key.shape[1],
+            )
+            if attn_bias_t.shape != expected_shape:
+                raise ValueError(
+                    f"Invalid shape for attention bias: {attn_bias_t.shape} (expected {expected_shape})\n"
+                    f"  query.shape: {self.query.shape}\n"
+                    f"  key.shape  : {self.key.shape}\n"
+                    f"  value.shape: {self.value.shape}"
+                )
+        if isinstance(self.attn_bias, BlockDiagonalMask):
+            if any(x.shape[0] != 1 for x in qkv):
+                raise ValueError(
+                    f"Expected batch_size=1 when using block-diagonal bias\n"
+                    f"  query.shape: {self.query.shape}\n"
+                    f"  key.shape  : {self.key.shape}\n"
+                    f"  value.shape: {self.value.shape}"
+                )
+        if self.p < 0.0 or self.p > 1.0:
+            raise ValueError(f"Invalid dropout probability: p={self.p}")
+        # Check that shapes match between inputs
+        B, Mq = self.query.shape[:2]
+        K = self.query.shape[-1]
+        B, Mkv = self.key.shape[:2]
+        Kv = self.value.shape[-1]
+        quantized_kv_cache = self.value.dtype == torch.int32
+        key_embed_dim = Kv if quantized_kv_cache else K
+
+        valid_shapes = True
+        if self.query.ndim == 3:  # BMK
+            valid_shapes = (
+                self.query.shape == (B, Mq, K)
+                and self.key.shape == (B, Mkv, K)
+                and self.value.shape == (B, Mkv, Kv)
+            )
+        H = self.query.shape[-2]
+        if self.query.ndim == 4:  # BMHK
+            valid_shapes = (
+                self.query.shape == (B, Mq, H, K)
+                and self.key.shape == (B, Mkv, H, key_embed_dim)
+                and self.value.shape == (B, Mkv, H, Kv)
+            )
+        G = self.query.shape[2]
+        if self.query.ndim == 5:  # BMNHK
+            valid_shapes = (
+                self.query.shape == (B, Mq, G, H, K)
+                and self.key.shape == (B, Mkv, G, H, key_embed_dim)
+                and self.value.shape == (B, Mkv, G, H, Kv)
+            )
+        if not valid_shapes:
+            raise ValueError(
+                f"Incompatible shapes for attention inputs:\n"
+                f"  query.shape: {self.query.shape}\n"
+                f"  key.shape  : {self.key.shape}\n"
+                f"  value.shape: {self.value.shape}\n"
+                "HINT: We don't support broadcasting, please use `expand` "
+                "yourself before calling `memory_efficient_attention` if you need to"
+            )
+
+    def get_output_dtype(self) -> torch.dtype:
+        if self.output_dtype is None:
+            if self.is_partial and self.query.dtype is not torch.float64:
+                return torch.float32
+            return self.query.dtype
+        return self.output_dtype
+
+    @property
+    def nbytes(self) -> int:
+        """
+        Number of bytes in the input, not counting the attention bias.
+        """
+        return sum(
+            x.untyped_storage().nbytes() for x in [self.query, self.key, self.value]
+        )
+
+
+@dataclass
+class Context:
+    lse: torch.Tensor
+    out: torch.Tensor
+    # NOTE: If `rng_state` is set, `op_bw` should be set as well
+    # as the randomness is backend-dependant
+    op_bw: Optional[Type["AttentionBwOpBase"]] = None
+    rng_state: Optional[Any] = None
+    qkv_share_storage: bool = False
+
+    def get_padded_lse(self, pad_to: int, force_pad_inf: bool = False) -> torch.Tensor:
+        pad_amount = (pad_to - (self.lse.shape[2] % pad_to)) % pad_to
+        lse = self.lse
+        if pad_amount > 0:
+            if force_pad_inf:
+                lse = lse[:, :, : self.out.shape[1]]
+                pad_amount = (pad_to - (lse.shape[2] % pad_to)) % pad_to
+            lse = torch.nn.functional.pad(lse, [0, pad_amount], value=math.inf)
+        elif force_pad_inf and self.out.shape[1] != lse.shape[2]:
+            lse[:, :, self.out.shape[1] :].fill_(math.inf)
+        return lse
+
+
+@dataclass
+class Gradients:
+    dq: torch.Tensor
+    dk: torch.Tensor
+    dv: torch.Tensor
+    # bias gradient. None if there is no tensor bias or if it doesn't require grad
+    db: Optional[torch.Tensor] = None
+
+
+class AttentionOpBase(BaseOperator):
+    """Base class for any attention operator in xFormers
+
+    See:
+
+    - :attr:`xformers.ops.fmha.cutlass.FwOp`
+    - :attr:`xformers.ops.fmha.cutlass.BwOp`
+    - :attr:`xformers.ops.fmha.flash.FwOp`
+    - :attr:`xformers.ops.fmha.flash.BwOp`
+    - :attr:`xformers.ops.fmha.triton.FwOp`
+    - :attr:`xformers.ops.fmha.triton.BwOp`
+    """
+
+    OPERATOR: Any  # pyre-ignore[13]
+    SUPPORTED_DEVICES: Set[str]  # pyre-ignore[13]
+    CUDA_MINIMUM_COMPUTE_CAPABILITY: Tuple[int, int] = (5, 0)
+    CUDA_MAXIMUM_COMPUTE_CAPABILITY: Optional[Tuple[int, int]] = None
+    SUPPORTED_DTYPES: Set[torch.dtype]  # pyre-ignore[13]
+    SUPPORTED_MAX_K: float  # pyre-ignore[13]
+    SUPPORTED_MIN_K: int = 0
+    SUPPORTED_ATTN_BIAS_TYPES: Iterable[Any] = (type(None),)
+    SUPPORTS_DROPOUT: bool  # pyre-ignore[13]
+    SUPPORTS_CUSTOM_SCALE: bool = False
+    SUPPORTS_DIFFERENT_VALUE_EMBED: bool = False
+    SUPPORTS_OUTPUT_DTYPE: bool = False
+    SUPPORTS_PARTIAL: bool = False
+    IS_DETERMINISTIC: bool = True
+    SUPPORTS_BMGHK: bool = False
+    NAME: str  # pyre-ignore[13]
+    OPERATOR_CATEGORY = "memory_efficient_attention"
+    # Format for the LSE computed in the FW pass, and accepted in the BW pass,
+    # for BlockDiagonalMask and children.
+    # When using a varlen bias, both the FW and BW operators must have the
+    # same value for `VARLEN_LSE_PACKED`
+    VARLEN_LSE_PACKED: bool = True
+
+    _TEST_BATCH_SIZES: List[int] = [1, 300]
+    _TEST_K: List[int] = [32, 128]
+
+    @classmethod
+    def supports(cls, d: Inputs) -> bool:
+        return not cls.not_supported_reasons(d)
+
+    @classmethod
+    def shape_not_supported_reasons(
+        cls, Mq: int, Mkv: int, K: int, Kv: int
+    ) -> List[str]:
+        reasons = []
+        if not cls.SUPPORTS_DIFFERENT_VALUE_EMBED and K != Kv:
+            reasons.append("query.shape[-1] != value.shape[-1]")
+        if max(K, Kv) > cls.SUPPORTED_MAX_K:
+            reasons.append(
+                f"max(query.shape[-1], value.shape[-1]) > {cls.SUPPORTED_MAX_K}"
+            )
+        if min(K, Kv) < cls.SUPPORTED_MIN_K:
+            reasons.append(
+                f"min(query.shape[-1], value.shape[-1]) < {cls.SUPPORTED_MIN_K}"
+            )
+        return reasons
+
+    @classmethod
+    def not_supported_reasons(cls, d: Inputs) -> List[str]:  # noqa: C901
+        """
+        Returns a list of reasons why this is not supported.
+        The kernel can run these inputs only if the returned list is empty
+        """
+        query_shape = d.query.shape
+        reasons = cls.shape_not_supported_reasons(
+            Mq=query_shape[1],
+            Mkv=d.key.shape[1],
+            K=query_shape[-1],
+            Kv=query_shape[-1] if d.value.dtype == torch.int32 else d.value.shape[-1],
+        )
+        device_type = d.query.device.type
+        dtype = d.query.dtype
+        if device_type not in cls.SUPPORTED_DEVICES:
+            reasons.append(f"device={device_type} (supported: {cls.SUPPORTED_DEVICES})")
+        if (
+            device_type == "cuda"
+            and not _built_with_cuda
+            and (torch.version.hip is None)
+        ):
+            reasons.append("xFormers wasn't build with CUDA support")
+        if device_type == "cuda" and (torch.version.hip is None):
+            device_capability = torch.cuda.get_device_capability(d.device)
+            if device_capability < cls.CUDA_MINIMUM_COMPUTE_CAPABILITY:
+                reasons.append(
+                    f"requires device with capability >= {cls.CUDA_MINIMUM_COMPUTE_CAPABILITY} "
+                    f"but your GPU has capability {device_capability} (too old)"
+                )
+            elif (
+                cls.CUDA_MAXIMUM_COMPUTE_CAPABILITY is not None
+                and device_capability > cls.CUDA_MAXIMUM_COMPUTE_CAPABILITY
+            ):
+                reasons.append(
+                    f"requires device with capability <= {cls.CUDA_MAXIMUM_COMPUTE_CAPABILITY} "
+                    f"but your GPU has capability {device_capability} (too new)"
+                )
+        if dtype not in cls.SUPPORTED_DTYPES:
+            reasons.append(f"dtype={dtype} (supported: {cls.SUPPORTED_DTYPES})")
+        if type(d.attn_bias) not in cls.SUPPORTED_ATTN_BIAS_TYPES:
+            reasons.append(f"attn_bias type is {type(d.attn_bias)}")
+        if not cls.SUPPORTS_OUTPUT_DTYPE:
+            if d.output_dtype is not None and d.output_dtype is not dtype:
+                reasons.append("Custom output dtype not supported")
+        if d.is_partial and not cls.SUPPORTS_PARTIAL:
+            reasons.append("Partial attention not supported")
+        if (d.p != 0.0) and not cls.SUPPORTS_DROPOUT:
+            reasons.append("dropout > 0.0")
+        if d.scale is not None and not cls.SUPPORTS_CUSTOM_SCALE:
+            reasons.append("has custom scale")
+        # bfloat16 is only supported on A100+ and MTIA
+        # ... although the kernels can still run and give the
+        # correct result
+        supports_bf16 = (
+            device_type.startswith("cuda")
+            and torch.cuda.get_device_capability(d.query.device)[0] >= 8
+        ) or device_type.startswith("mtia")
+        if dtype is torch.bfloat16 and not supports_bf16:
+            reasons.append("bf16 is only supported on A100+ GPUs and MTIA")
+        if not cls.is_available():
+            reasons.append(
+                "operator wasn't built - see `python -m xformers.info` for more info"
+            )
+        if not cls.IS_DETERMINISTIC and torch.are_deterministic_algorithms_enabled():
+            reasons.append(
+                "operator is non-deterministic, but `torch.use_deterministic_algorithms` is set"
+            )
+        if not cls.SUPPORTS_BMGHK and d.query.ndim == 5:
+            reasons.append("operator does not support BMGHK format")
+        return reasons
+
+
+class AttentionFwOpBase(AttentionOpBase):
+    ERROR_ATOL: Mapping[torch.dtype, float] = {
+        torch.float: 3e-4,
+        torch.half: 4e-3,
+        torch.bfloat16: 2e-2,
+    }
+    ERROR_RTOL: Mapping[torch.dtype, float] = {
+        torch.float: 2e-5,
+        torch.half: 4e-4,
+        torch.bfloat16: 5e-3,
+    }
+
+    @classmethod
+    def apply(
+        cls, inp: Inputs, needs_gradient: bool
+    ) -> Tuple[torch.Tensor, Optional[Context]]:
+        raise NotImplementedError()
+
+
+class AttentionBwOpBase(AttentionOpBase):
+    # NOTE on tolerances: These are tested for `scales => (1/32)**0.5`
+    # In the BW pass, imprecisions accumulate in the Q@K.T recalculation
+    # These imprecisions are multiplied by the `scale` and then exponentiated
+    # So if the scale is too high, we get a lot of errors
+
+    ERROR_ATOL: Mapping[torch.dtype, float] = {
+        torch.float: 9e-4,
+        torch.half: 0.2,
+        torch.bfloat16: 0.9,
+    }
+    ERROR_RTOL: Mapping[torch.dtype, float] = {
+        torch.float: 1e-4,
+        torch.half: 2e-2,
+        torch.bfloat16: 0.1,
+    }
+    SUPPORTS_ATTN_BIAS_GRAD = False
+    SUPPORTS_PARTIAL = True
+
+    @classmethod
+    def not_supported_reasons(cls, d: Inputs) -> List[str]:
+        reasons = super(AttentionBwOpBase, cls).not_supported_reasons(d)
+        if (
+            isinstance(d.attn_bias, torch.Tensor)
+            and d.attn_bias.requires_grad
+            and not cls.SUPPORTS_ATTN_BIAS_GRAD
+        ):
+            reasons.append(
+                "Computing the bias gradient is not supported (attn_bias.requires_grad = True)"
+            )
+
+        return reasons
+
+    @classmethod
+    def apply(cls, ctx: Context, inp: Inputs, grad: torch.Tensor) -> Gradients:
+        raise NotImplementedError()
+
+
+AttentionOp = Tuple[
+    Optional[Type[AttentionFwOpBase]], Optional[Type[AttentionBwOpBase]]
+]
+
+
+def bmk2bmhk(tensor, num_heads: int) -> torch.Tensor:
+    if tensor.ndim == 4:
+        return tensor
+    return tensor.reshape(
+        [tensor.shape[0] // num_heads, num_heads, tensor.shape[1], tensor.shape[2]]
+    ).permute((0, 2, 1, 3))
+
+
+def check_lastdim_alignment_stride1(
+    reasons: List[str], name: str, x: torch.Tensor, alignment: int
+) -> None:
+    if x.shape[-1] % alignment != 0:
+        reasons.append(f"{name}.shape[-1] % {alignment} != 0")
+    elif x.stride(-2) % alignment != 0:
+        reasons.append(
+            f"{name}.stride(-2) % {alignment} != 0 ({name}.stride() = {x.stride()})"
+        )
+    # We can have stride=0 sometimes if dimension=1
+    if x.stride(-1) > 1:
+        reasons.append(
+            f"{name}.stride(-1) > 1 ({name}.stride() = {x.stride()}) - you should call `.contiguous()` on the input"
+        )