pytorch - ✅(Solved) Fix torch.export: nn.MultiheadAttention (need_weights=False) exported on CPU fails on CUDA (permute→view after SDPA) [1 pull requests, 1 comments, 2 participants]

pytorch2026-04-03 19:10:16

ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

GitHub issue URL

Helpful · Quick feedback

GitHub stats

pytorch/pytorch#179287•Fetched 2026-04-08 02:43:40

View on GitHub

Comments

Participants

Timeline

Reactions

Author

Participants

Assignees

Timeline (top)

mentioned ×8subscribed ×8labeled ×3assigned ×2

Error Message

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

Fix Action

Fixed

Fixed by PR: fix(export): add SDPA decomposition for contiguous output (https://github.com/pytorch/pytorch/pull/179376)

PR fix notes

PR #179376: fix(export): add SDPA decomposition for contiguous output

Repository: pytorch/pytorch
Author: WeishuZ
State: open | merged: False
Link: https://github.com/pytorch/pytorch/pull/179376

Description (problem / solution / changelog)

Summary

When exporting nn.MultiheadAttention with need_weights=False and running on CUDA, the exported graph contains permute->view without contiguous() in between, causing a view error.

Problem

The exported graph shows:

scaled_dot_product_attention = ...
permute = torch.ops.aten.permute.default(...)
view_8 = torch.ops.aten.view.default(...) # fails on CUDA

The contiguous() call is missing between permute and view.

Solution

Added a decomposition for aten.scaled_dot_product_attention that ensures the output tensor is contiguous. This maintains compatibility with the permute->view pattern in MultiheadAttention when running on CUDA.

Changes

torch/_decomp/decompositions.py: Added decomposition for scaled_dot_product_attention
torch/_decomp/init.py: Registered the new decomposition

Fixes #179287

Changed files

torch/_decomp/__init__.py (modified, +1/-0)
torch/_decomp/decompositions.py (modified, +74/-0)

Code Example

scaled_dot_product_attention = torch.ops.aten.scaled_dot_product_attention.default(...)
permute = torch.ops.aten.permute.default(scaled_dot_product_attention, [2, 0, 1, 3])
view_8 = torch.ops.aten.view.default(permute, [...])

---

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

---

import torch
from torch import nn


class ReproModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = nn.MultiheadAttention(embed_dim=16, num_heads=1, batch_first=True, dropout=0.0)

    def forward(self, inputs: torch.Tensor, mask: torch.Tensor):
        outputs, _ = self.attention(inputs, inputs, inputs, key_padding_mask=~mask, need_weights=False)
        return outputs


assert torch.cuda.is_available()

cpu_inputs = torch.randn(2, 7, 16)
cpu_mask = torch.tensor([
    [True, True, True, True, True, False, False],
    [True, True, True, False, False, False, False],
])

with torch.no_grad():
    exported_program = torch.export.export(
        ReproModel().eval().cpu(),
        args=(cpu_inputs, cpu_mask),
        strict=False,
    )

print(exported_program.graph_module.code)
exported_module = exported_program.module()

with torch.no_grad():
    exported_module(cpu_inputs, cpu_mask)
print("CPU run: success")

exported_module = exported_module.cuda()
with torch.no_grad():
    exported_module(cpu_inputs.cuda(), cpu_mask.cuda())  # fails here
print("CUDA run: success")

---

def forward(self, p_attention_in_proj_weight, p_attention_in_proj_bias, p_attention_out_proj_weight, p_attention_out_proj_bias, inputs, mask):
    bitwise_not = torch.ops.aten.bitwise_not.default(mask);  mask = None
    zeros_like = torch.ops.aten.zeros_like.default(bitwise_not, dtype = torch.float32, pin_memory = False)
    masked_fill_ = torch.ops.aten.masked_fill_.Scalar(zeros_like, bitwise_not, -inf);  zeros_like = bitwise_not = None
    transpose = torch.ops.aten.transpose.int(inputs, 1, 0);  inputs = None
    linear = torch.ops.aten.linear.default(transpose, p_attention_in_proj_weight, p_attention_in_proj_bias);  transpose = p_attention_in_proj_weight = p_attention_in_proj_bias = None
    unflatten = torch.ops.aten.unflatten.int(linear, -1, [3, 16]);  linear = None
    unsqueeze = torch.ops.aten.unsqueeze.default(unflatten, 0);  unflatten = None
    transpose_1 = torch.ops.aten.transpose.int(unsqueeze, 0, -2);  unsqueeze = None
    squeeze = torch.ops.aten.squeeze.dim(transpose_1, -2);  transpose_1 = None
    contiguous = torch.ops.aten.contiguous.default(squeeze);  squeeze = None
    select = torch.ops.aten.select.int(contiguous, 0, 0)
    select_1 = torch.ops.aten.select.int(contiguous, 0, 1)
    select_2 = torch.ops.aten.select.int(contiguous, 0, 2);  contiguous = None
    view = torch.ops.aten.view.default(select, [7, 2, 16]);  select = None
    transpose_2 = torch.ops.aten.transpose.int(view, 0, 1);  view = None
    view_1 = torch.ops.aten.view.default(select_1, [7, 2, 16]);  select_1 = None
    transpose_3 = torch.ops.aten.transpose.int(view_1, 0, 1);  view_1 = None
    view_2 = torch.ops.aten.view.default(select_2, [7, 2, 16]);  select_2 = None
    transpose_4 = torch.ops.aten.transpose.int(view_2, 0, 1);  view_2 = None
    view_3 = torch.ops.aten.view.default(masked_fill_, [2, 1, 1, 7]);  masked_fill_ = None
    expand = torch.ops.aten.expand.default(view_3, [-1, 1, -1, -1]);  view_3 = None
    reshape = torch.ops.aten.reshape.default(expand, [2, 1, 7]);  expand = None
    view_4 = torch.ops.aten.view.default(reshape, [2, 1, -1, 7]);  reshape = None
    view_5 = torch.ops.aten.view.default(transpose_2, [2, 1, 7, 16]);  transpose_2 = None
    view_6 = torch.ops.aten.view.default(transpose_3, [2, 1, 7, 16]);  transpose_3 = None
    view_7 = torch.ops.aten.view.default(transpose_4, [2, 1, 7, 16]);  transpose_4 = None
    scaled_dot_product_attention = torch.ops.aten.scaled_dot_product_attention.default(view_5, view_6, view_7, view_4);  view_5 = view_6 = view_7 = view_4 = None
    permute = torch.ops.aten.permute.default(scaled_dot_product_attention, [2, 0, 1, 3]);  scaled_dot_product_attention = None
    view_8 = torch.ops.aten.view.default(permute, [14, 16]);  permute = None
    linear_1 = torch.ops.aten.linear.default(view_8, p_attention_out_proj_weight, p_attention_out_proj_bias);  view_8 = p_attention_out_proj_weight = p_attention_out_proj_bias = None
    view_9 = torch.ops.aten.view.default(linear_1, [7, 2, 16]);  linear_1 = None
    transpose_5 = torch.ops.aten.transpose.int(view_9, 1, 0);  view_9 = None
    return (transpose_5,)
    
CPU run: success
Traceback (most recent call last):
  File "/tmp/repro.py", line 39, in <module>
    exported_module(cpu_inputs.cuda(), cpu_mask.cuda())  # fails here
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 949, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 461, in __call__
    raise e
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 447, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1885, in _call_impl
    return inner()
           ^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1833, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.9", line 39, in forward
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/_ops.py", line 865, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

torch.export produces an ExportedProgram for nn.MultiheadAttention where the exported graph contains aten.permute followed directly by aten.view on the SDPA output, with no aten.contiguous in between.

Tested with torch 2.11.0+cu130 and Python 3.11.

The exported graph contains:

scaled_dot_product_attention = torch.ops.aten.scaled_dot_product_attention.default(...)
permute = torch.ops.aten.permute.default(scaled_dot_product_attention, [2, 0, 1, 3])
view_8 = torch.ops.aten.view.default(permute, [...])

There is no contiguous between the permute and view.

Running the exported module succeeds on CPU, but fails after being moved to a CUDA device with:

RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

Example to reproduce:

import torch
from torch import nn


class ReproModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.attention = nn.MultiheadAttention(embed_dim=16, num_heads=1, batch_first=True, dropout=0.0)

    def forward(self, inputs: torch.Tensor, mask: torch.Tensor):
        outputs, _ = self.attention(inputs, inputs, inputs, key_padding_mask=~mask, need_weights=False)
        return outputs


assert torch.cuda.is_available()

cpu_inputs = torch.randn(2, 7, 16)
cpu_mask = torch.tensor([
    [True, True, True, True, True, False, False],
    [True, True, True, False, False, False, False],
])

with torch.no_grad():
    exported_program = torch.export.export(
        ReproModel().eval().cpu(),
        args=(cpu_inputs, cpu_mask),
        strict=False,
    )

print(exported_program.graph_module.code)
exported_module = exported_program.module()

with torch.no_grad():
    exported_module(cpu_inputs, cpu_mask)
print("CPU run: success")

exported_module = exported_module.cuda()
with torch.no_grad():
    exported_module(cpu_inputs.cuda(), cpu_mask.cuda())  # fails here
print("CUDA run: success")

Sample output:

def forward(self, p_attention_in_proj_weight, p_attention_in_proj_bias, p_attention_out_proj_weight, p_attention_out_proj_bias, inputs, mask):
    bitwise_not = torch.ops.aten.bitwise_not.default(mask);  mask = None
    zeros_like = torch.ops.aten.zeros_like.default(bitwise_not, dtype = torch.float32, pin_memory = False)
    masked_fill_ = torch.ops.aten.masked_fill_.Scalar(zeros_like, bitwise_not, -inf);  zeros_like = bitwise_not = None
    transpose = torch.ops.aten.transpose.int(inputs, 1, 0);  inputs = None
    linear = torch.ops.aten.linear.default(transpose, p_attention_in_proj_weight, p_attention_in_proj_bias);  transpose = p_attention_in_proj_weight = p_attention_in_proj_bias = None
    unflatten = torch.ops.aten.unflatten.int(linear, -1, [3, 16]);  linear = None
    unsqueeze = torch.ops.aten.unsqueeze.default(unflatten, 0);  unflatten = None
    transpose_1 = torch.ops.aten.transpose.int(unsqueeze, 0, -2);  unsqueeze = None
    squeeze = torch.ops.aten.squeeze.dim(transpose_1, -2);  transpose_1 = None
    contiguous = torch.ops.aten.contiguous.default(squeeze);  squeeze = None
    select = torch.ops.aten.select.int(contiguous, 0, 0)
    select_1 = torch.ops.aten.select.int(contiguous, 0, 1)
    select_2 = torch.ops.aten.select.int(contiguous, 0, 2);  contiguous = None
    view = torch.ops.aten.view.default(select, [7, 2, 16]);  select = None
    transpose_2 = torch.ops.aten.transpose.int(view, 0, 1);  view = None
    view_1 = torch.ops.aten.view.default(select_1, [7, 2, 16]);  select_1 = None
    transpose_3 = torch.ops.aten.transpose.int(view_1, 0, 1);  view_1 = None
    view_2 = torch.ops.aten.view.default(select_2, [7, 2, 16]);  select_2 = None
    transpose_4 = torch.ops.aten.transpose.int(view_2, 0, 1);  view_2 = None
    view_3 = torch.ops.aten.view.default(masked_fill_, [2, 1, 1, 7]);  masked_fill_ = None
    expand = torch.ops.aten.expand.default(view_3, [-1, 1, -1, -1]);  view_3 = None
    reshape = torch.ops.aten.reshape.default(expand, [2, 1, 7]);  expand = None
    view_4 = torch.ops.aten.view.default(reshape, [2, 1, -1, 7]);  reshape = None
    view_5 = torch.ops.aten.view.default(transpose_2, [2, 1, 7, 16]);  transpose_2 = None
    view_6 = torch.ops.aten.view.default(transpose_3, [2, 1, 7, 16]);  transpose_3 = None
    view_7 = torch.ops.aten.view.default(transpose_4, [2, 1, 7, 16]);  transpose_4 = None
    scaled_dot_product_attention = torch.ops.aten.scaled_dot_product_attention.default(view_5, view_6, view_7, view_4);  view_5 = view_6 = view_7 = view_4 = None
    permute = torch.ops.aten.permute.default(scaled_dot_product_attention, [2, 0, 1, 3]);  scaled_dot_product_attention = None
    view_8 = torch.ops.aten.view.default(permute, [14, 16]);  permute = None
    linear_1 = torch.ops.aten.linear.default(view_8, p_attention_out_proj_weight, p_attention_out_proj_bias);  view_8 = p_attention_out_proj_weight = p_attention_out_proj_bias = None
    view_9 = torch.ops.aten.view.default(linear_1, [7, 2, 16]);  linear_1 = None
    transpose_5 = torch.ops.aten.transpose.int(view_9, 1, 0);  view_9 = None
    return (transpose_5,)
    
CPU run: success
Traceback (most recent call last):
  File "/tmp/repro.py", line 39, in <module>
    exported_module(cpu_inputs.cuda(), cpu_mask.cuda())  # fails here
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 949, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 461, in __call__
    raise e
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/fx/graph_module.py", line 447, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1779, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1885, in _call_impl
    return inner()
           ^^^^^^^
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1833, in inner
    result = forward_call(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<eval_with_key>.9", line 39, in forward
  File "/home/ansible/Workspace/stonks/.venv/lib/python3.11/site-packages/torch/_ops.py", line 865, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: view size is not compatible with input tensor's size and stride (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead.

Versions

Collecting environment information... PyTorch version: 2.11.0+cu130 Is debug build: False CUDA used to build PyTorch: 13.0 ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.4 LTS (x86_64) GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04.1) 13.3.0 Clang version: Could not collect CMake version: Could not collect Libc version: glibc-2.39

Python version: 3.11.11 (main, Dec 19 2024, 14:33:27) [Clang 18.1.8 ] (64-bit runtime) Python platform: Linux-6.8.0-106-generic-x86_64-with-glibc2.39 Is CUDA available: True CUDA runtime version: Could not collect CUDA_MODULE_LOADING set to: GPU models and configuration: GPU 0: NVIDIA RTX 4000 SFF Ada Generation Nvidia driver version: 580.126.20 cuDNN version: Could not collect Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True Caching allocator config: N/A

CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 20 On-line CPU(s) list: 0-19 Vendor ID: GenuineIntel Model name: 13th Gen Intel(R) Core(TM) i5-13500 CPU family: 6 Model: 191 Thread(s) per core: 2 Core(s) per socket: 14 Socket(s): 1 Stepping: 2 CPU(s) scaling MHz: 74% CPU max MHz: 4800.0000 CPU min MHz: 800.0000 BogoMIPS: 4992.00

Virtualization: VT-x L1d cache: 544 KiB (14 instances) L1i cache: 704 KiB (14 instances) L2 cache: 11.5 MiB (8 instances) L3 cache: 24 MiB (1 instance) NUMA node(s): 1 NUMA node0 CPU(s): 0-19

cc @chauhang @penguinwu @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

extent analysis

TL;DR

Inserting a contiguous operation between aten.permute and aten.view in the exported graph should resolve the compatibility issue on CUDA devices.

Guidance

Identify the aten.permute and aten.view operations in the exported graph and insert aten.contiguous between them to ensure the tensor is contiguous in memory before applying the view operation.
Verify that the modified graph runs successfully on both CPU and CUDA devices.
Consider using torch.reshape instead of torch.view if the issue persists, as suggested by the error message.
Review the PyTorch documentation for torch.export and torch.fx to understand the graph manipulation and optimization process.

Example

# Insert aten.contiguous between permute and view
permute = torch.ops.aten.permute.default(scaled_dot_product_attention, [2, 0, 1, 3])
contiguous = torch.ops.aten.contiguous.default(permute)
view_8 = torch.ops.aten.view.default(contiguous, [14, 16])

Notes

The issue seems to be related to the memory layout of the tensor after the aten.permute operation, which is not compatible with the subsequent aten.view operation on CUDA devices. Inserting aten.contiguous should ensure the tensor is properly aligned in memory.

Recommendation

Apply the workaround by inserting aten.contiguous between aten.permute and aten.view to resolve the compatibility issue on CUDA devices.

Vote matrix · Quick signals

Works

Did the solution work? Tap to confirm.

Easy Fix

Was it a quick fix?

Time Saver

Did it save you time?

Blocking

Was it severely blocking?

Common Issue

Are others likely hitting this too?

Flaky / Intermittent

Is it intermittent?

Verified / Reproducible

Can you reproduce it reliably?

#api #ssr #installation #tensor shape #autograd error

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Data

Security

Network

Code

UI/UX

Text

System

Multimedia

Protocol

API

Engineering

pytorch - ✅(Solved) Fix torch.export: nn.MultiheadAttention (need_weights=False) exported on CPU fails on CUDA (permute→view after SDPA) [1 pull requests, 1 comments, 2 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fixed

PR fix notes

PR #179376: fix(export): add SDPA decomposition for contiguous output

Description (problem / solution / changelog)

Summary

Problem

Solution

Changes

Changed files

Code Example

🐛 Describe the bug

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

TRENDING

pytorch - ✅(Solved) Fix torch.export: nn.MultiheadAttention (need_weights=False) exported on CPU fails on CUDA (permute→view after SDPA) [1 pull requests, 1 comments, 2 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fixed

PR fix notes

PR #179376: fix(export): add SDPA decomposition for contiguous output

Description (problem / solution / changelog)

Summary

Problem

Solution

Changes

Changed files

Code Example

🐛 Describe the bug

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

RELATED_DISCOVERY

TRENDING