pytorch - 💡(How to fix) Fix Inductor failure with `view_as_complex` in fused Multi-Head Attention subgraph [1 participants]

import torch
import torch.nn as nn

class InductorBugRepro(nn.Module):
    def __init__(self, d_model):
        super().__init__()
        self.w_q = nn.Parameter(torch.randn(d_model, d_model))
        self.w_k = nn.Parameter(torch.randn(d_model, d_model))

    def forward(self, x, rope):
        B, L, _ = x.shape
        H, D = 16, 64
        
        # 1. Linear projection + Transpose (Standard MHSA pattern)
        q = (x @ self.w_q.T).view(B, L, H, D).transpose(1, 2)
        k = (x @ self.w_k.T).view(B, L, H, D).transpose(1, 2)
        
        # 2. Parallel complex op chains (RoPE pattern)
        # Each chain uses .contiguous() before view_as_complex
        q_rot = torch.view_as_real(torch.view_as_complex(q.view(*q.shape[:-1], -1, 2).contiguous()) * rope)
        k_rot = torch.view_as_real(torch.view_as_complex(k.view(*k.shape[:-1], -1, 2).contiguous()) * rope)
        
        # 3. Matmul consumer forces fusion of the RoPE outputs
        return torch.matmul(q_rot.view(*q.shape), k_rot.view(*k.shape).transpose(-1, -2))

def run():
    device = "cuda"
    B, H, L, D = 1, 16, 1, 64
    d_model = H * D
    
    model = InductorBugRepro(d_model).to(device)
    compiled_model = torch.compile(model)
    
    x = torch.randn(B, L, d_model, device=device)
    rope = torch.randn(L, D // 2, dtype=torch.complex64, device=device)
    
    print(f"Running reproduction on {device}...")
    
    # Passes
    _ = model(x, rope)
    print("Eager mode: Success")
    
    # Fails
    try:
        _ = compiled_model(x, rope)
        print("Compiled mode: Success")
    except Exception as e:
        print(f"Compiled mode: FAILED\n\nError:\n{e}")

if __name__ == "__main__":
    run()

---

Running reproduction on cuda...
Eager mode: Success
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0] failed while attempting to run meta for aten.view_as_complex.default
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0] Traceback (most recent call last):
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]   File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2823, in _dispatch_impl
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]     r = func(*args, **kwargs)
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]         ^^^^^^^^^^^^^^^^^^^^^
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]   File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 819, in __call__
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]     return self._op(*args, **kwargs)
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0]            ^^^^^^^^^^^^^^^^^^^^^^^^^
E0310 03:16:58.968000 463 torch/_subclasses/fake_tensor.py:2827] [1/0] RuntimeError: Tensor must have a stride divisible by 2 for all but last dimension
Compiled mode: FAILED

Error:
backend='inductor' raised:
RuntimeError: Tensor must have a stride divisible by 2 for all but last dimension

Set TORCHDYNAMO_VERBOSE=1 for the internal stack trace (please do this especially if you're reporting a bug to PyTorch). For even more developer context, set TORCH_LOGS="+dynamo"

/usr/local/lib/python3.12/dist-packages/torch/autograd/graph.py:865: UserWarning: Error detected in ViewAsRealBackward0. Traceback of forward call that caused the error:
  File "/tmp/ipykernel_463/4071400647.py", line 21, in forward
    k_rot = torch.view_as_real(torch.view_as_complex(k.view(*k.shape[:-1], -1, 2).contiguous()) * rope)
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass