[Issue]: `invalid argument for fmha_fwd` with torch.Compile in gfx950

[jjuvonen-amd]

Problem Description

Aiter FAv3 forward throws RuntimeError: invalid argument for fmha_fwd when model is torch.Compiled and attention inputs are transposed from BHSD to BSHD format before attention function.

Error occurs only with Compile and if .transpose(1,2).contiguous() is applied for q,k,v inputs. Also, error happens on gfx950, but not on gfx942 (although differences between eager vs. compile FAv3 fwd outputs is suspiciously high >5e-2 on gfx942).

Operating System

Ubuntu 24.04.2 LTS (Noble Numbat)

CPU

AMD EPYC 9575F 64-Core Processor

GPU

AMD Instinct MI355X

ROCm Version

ROCm 7.0.0rc20250820

ROCm Component

No response

Steps to Reproduce

Running pytest test_aiter_fa_compile.py runs 6 tests:

• test_aiter_attention[Eager-BSHD order] Passes
• test_aiter_attention[Eager-BHSD order] Passes
• test_aiter_attention[Compiled-BSHD order] Passes
• test_aiter_attention[Compiled-BHSD order] Fails on gfx950 invalid argument for fmha_fwd
• test_compare_eager_vs_compiled[BSHD order] Passes
• test_compare_eager_vs_compiled[BHSD order] Fails on gfx950 invalid argument for fmha_fwd

test_aiter_fa_compile.py:

import pytest
import torch
import numpy as np
import aiter

class SimpleAiterAttentionModel(torch.nn.Module):
    """Simple model with just flash attention."""
    
    def __init__(self, bhsd_order=False):
        super().__init__()
        self.flash_attn_func = aiter.flash_attn_func
        self.bhsd_order = bhsd_order # If True, expects input in (B, H, S, D) order
            
    def forward(self, q, k, v):
        if self.bhsd_order:
            # Convert from (B, H, S, D) to (B, S, H, D) for flash attention
            q = q.transpose(1,2).contiguous()
            k = k.transpose(1,2).contiguous()
            v = v.transpose(1,2).contiguous()

        output = self.flash_attn_func(
            q, k, v, 
            causal=False, 
            deterministic=False,
            return_lse=True # gfx950 requires this to use FAv3 fwd
        )[0]
        
        if self.bhsd_order:
            # Convert back to (B, H, S, D)
            output = output.transpose(1,2).contiguous()
        
        return output

def get_test_data(batch_size=1, seq_len=32768, heads=16, dim=128, bhsd_order=False):
    """Generate test data for attention models."""
    torch.manual_seed(42)

    shape = (batch_size, seq_len, heads, dim)
    if bhsd_order:
        shape = (batch_size, heads, seq_len, dim)
    
    def _generate_data():
        return torch.randn(shape, 
                           device='cuda', 
                           dtype=torch.bfloat16, 
                           requires_grad=False)
    q = _generate_data()
    k = _generate_data()
    v = _generate_data()
    return q, k, v

@pytest.mark.parametrize("bhsd_order", [False, True], ids=["BSHD order", "BHSD order"])
@pytest.mark.parametrize("use_compile", [False, True], ids=["Eager", "Compiled"])
def test_aiter_attention(bhsd_order, use_compile):
    """Test aiter flash attention with and without torch.compile."""
    q, k, v = get_test_data(bhsd_order=bhsd_order)
    
    # Create model and run forward pass
    model = SimpleAiterAttentionModel(bhsd_order=bhsd_order).eval()
    
    if use_compile:
        model = torch.compile(model, mode='default')

    with torch.inference_mode(True):
        output = model(q, k, v)
    
    # check for nans or infs or all zeros
    assert not torch.isnan(output).any(), "Output contains NaNs"
    assert not torch.isinf(output).any(), "Output contains Infs"
    assert not torch.all(output == 0), "Output is all zeros"

@pytest.mark.parametrize("bhsd_order", [False, True], ids=["BSHD order", "BHSD order"])
def test_compare_eager_vs_compiled(bhsd_order):
    """Compare outputs from eager and compiled execution."""
    q, k, v = get_test_data(bhsd_order=bhsd_order)
    
    # Get outputs from eager and compiled versions
    model = SimpleAiterAttentionModel(bhsd_order=bhsd_order).eval()
    with torch.inference_mode(True):
        eager_out = model(q.clone(), k.clone(), v.clone())

    model_compiled = torch.compile(model)
    with torch.inference_mode(True):
        compiled_out = model_compiled(q.clone(), k.clone(), v.clone())

    np.testing.assert_allclose(
        eager_out.to(torch.float32).cpu().numpy(),
        compiled_out.to(torch.float32).cpu().numpy(),
        rtol=1e-1, # generous tolerances due to bfloat16
        atol=5e-2,
        err_msg=(
            f"Outputs from eager and compiled model do not match. "
        )
    )

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

Torch version: 2.7.1+rocmsdk20250821
Aiter version: 0.1.5.dev157+g01330f6c6

[ZhangLirong-amd]

@jjuvonen-amd ,hi, can you give me your docker image name and more environment detail? I running on MI350, and cannot reproduce the error. Thanks.

test_aiter_fa_compile.py::test_aiter_attention[Eager-BSHD order] PASSED                                                                   [ 25%]
test_aiter_fa_compile.py::test_aiter_attention[Eager-BHSD order] PASSED                                                                   [ 50%]
test_aiter_fa_compile.py::test_aiter_attention[Compiled-BSHD order] PASSED                                                                [ 75%]
test_aiter_fa_compile.py::test_aiter_attention[Compiled-BHSD order] PASSED                                                                [100%]

[jjuvonen-amd]

Thanks @ZhangLirong-amd for looking into this. The docker image is not in registry, but the Dockerfile is attached here (as .txt file).
Build time is ~30min (installs triton, TheRock, and Aiter from source - pinned to commit hashes).

I've only tried this on MI355X and not in MI350 (sorry, don't have more info about the compute cluster)
Dockerfile.m355x-aiter01330f6c6.txt

[jjuvonen-amd]

The original invalid argument for fmha_fwd error has been resolved in the latest commit: 256bf677a.

Also, the difference between eager and compile outputs in test_compare_eager_vs_compiled[BHSD order] unit test is related to having .contiguous() called for aiter.flash_attn_func inputs. Torch.Compile doesn't seem to actually make tensors contiguous in that case. However, adding a graph break between contiguous() and flash_attention_func will fix this and make eager and Compile outputs similar.

This model will pass all unit tests and produce similar outputs in eager and compile modes:

@torch.compiler.disable()
def _noop(tensor):
    return tensor

class SimpleAiterAttentionModel(torch.nn.Module):
    """Simple model with just flash attention."""
    
    def __init__(self, bhsd_order=False):
        super().__init__()
        self.flash_attn_func = aiter.flash_attn_func
        self.bhsd_order = bhsd_order # If True, expects input in (B, H, S, D) order
            
    def forward(self, q, k, v):
        if self.bhsd_order:
            # Convert from (B, H, S, D) to (B, S, H, D) for flash attention
            q = q.transpose(1,2).contiguous()
            k = k.transpose(1,2).contiguous()
            v = v.transpose(1,2).contiguous()

            # without this, torch.compile doesn't seem to actually apply contiguous() - suspicion
            q = _noop(q)
            k = _noop(k)
            v = _noop(v)
        
        output = self.flash_attn_func(
            q, k, v, 
            causal=False, 
            deterministic=False,
            return_lse=True # gfx950 requires this to use FAv3 fwd
        )[0]
        
        if self.bhsd_order:
            # Convert back to (B, H, S, D)
            output = output.transpose(1,2)
            output = output.contiguous()
        
        return output

[ZhangLirong-amd]

@jjuvonen-amd , Hi, you are right, Torch.Compile doesn't seem to actually make tensors contiguous, and got the wrong stride for q,k,v. We thought it maybe a bug in torch.compile. and no better solution yet. You can avoid compile q = q.transpose(1,2).contiguous() as you mentioned before, and can also try torch 2.8. We got right stride in 2.8. We have asked the experts at torch.compile to check it together. I will tell you if we find a better way to solve this stride error in torch.compile

[jataylo]

@jjuvonen-amd, @ZhangLirong-amd I cannot reproduce this issue with tip of ROCm release/2.7 branch, test_aiter_attention[Compiled-BHSD order] passes as is, and I seem to be hitting some abort on final test, but not the invalid arg.

test.py::test_aiter_attention[Eager-BSHD order] PASSED                                                       [ 16%]
test.py::test_aiter_attention[Eager-BHSD order] PASSED                                                       [ 33%]
test.py::test_aiter_attention[Compiled-BSHD order] PASSED                                                    [ 50%]
test.py::test_aiter_attention[Compiled-BHSD order] PASSED                                                    [ 66%]
test.py::test_compare_eager_vs_compiled[BSHD order] PASSED                                                   [ 83%]
test.py::test_compare_eager_vs_compiled[BHSD order] Fatal Python error: Aborted

@jjuvonen-amd Could you give tip of this branch a try and see if the fundamental issue is resolved?

[jjuvonen-amd]

@jataylo, I tried this now with a very recent pytorch ROCm release/2.7 branch (commit: 59925f5) and aiter (commit: 256bf677a)

The original invalid argument for fmha_fwd error is gone but torch.compile output differs from eager when transpose+contiguous is used for aiter.FA inputs.

The only failing test is test_compare_eager_vs_compiled[BHSD order] - AssertionError:

E           AssertionError:
E           Not equal to tolerance rtol=0.1, atol=0.05
E           Outputs from eager and compiled model do not match.
E           Mismatched elements: 316 / 67108864 (0.000471%)
E           Max absolute difference among violations: 0.06677246
E           Max relative difference among violations: 41.745342

Perhaps this is caused by contiguous() not getting actually applied in torch.compile as discussed above with @ZhangLirong-amd

[amd-mghanimi]

I am using torch 2.8 and still got this error.

So, similar to @jjuvonen-amd I created a function with decorator for torch.compiler.disable() but I moved the transpose function there. it seems it fixes the issue in compile mode on MI300X where I tested this:

https://github.com/AMD-AIG-AIMA/torchtune/blob/da1594639e90c311e8d726fac971e2096defa388/torchtune/modules/attention_utils.py#L140-L153

@torch.compiler.disable()
def convert_to_bshd(tensor):
    return tensor.transpose(1, 2).contiguous()

[jataylo]

Seems like in 2.7 it reinterprets tensor as contiguous after the custom op not before, whereas 2.8 will correctly allocate as contiguous before the custom op calls. Contacted upstream pytorch inductor folks to see if they know what changed, if not I'll try to dive in and see what fixed it

I recommend to use 2.8 as a workaround until we can get this figured out.

[ZhangLirong-amd]

Thanks @jataylo help, we found vLLM added needs_fixed_stride_order to a large number of custom operators as a workaround. And upstream torch 2.7 is hard so we refer to vllm implement and make a PR in aiter to fix it #950 .

[avjves]

I feel like I'm still getting the same issue (or at least something very similar), tested with both pytorch 2.8 and pytorch 2.7 with the latest aiter commit (so with the PR #950 merged).

Here's a repro script:

import torch
import aiter

@torch.compiler.disable()
def contiguous_no_compile(tensor):
    return tensor.contiguous()

def attention_function_with_graph_break(q, k, v):
    q = contiguous_no_compile(q)
    k = contiguous_no_compile(k)
    v = contiguous_no_compile(v)
    out = aiter.flash_attn_func(q, k, v, causal=False)
    return out

def attention_function(q, k, v):
    q = q.contiguous()
    k = k.contiguous()
    v = v.contiguous()
    out = aiter.flash_attn_func(q, k, v, causal=False)
    return out

if __name__ == "__main__":

    q = torch.randn(29760, 2, 5, 128, device="cuda", dtype=torch.bfloat16)
    k = torch.randn(29760, 2, 5, 128, device="cuda", dtype=torch.bfloat16)
    v = torch.randn(29760, 2, 5, 128, device="cuda", dtype=torch.bfloat16)

    q = torch.permute(q, [1, 0, 2, 3])
    k = torch.permute(k, [1, 0, 2, 3])
    v = torch.permute(v, [1, 0, 2, 3])

    cont_q = q.contiguous()
    cont_k = k.contiguous()
    cont_v = v.contiguous()
    out_no_comp = aiter.flash_attn_func(cont_q, cont_k, cont_v, causal=False)

    #### Works
    compiled_with_graph_break = torch.compile(attention_function_with_graph_break)
    out_compiled_with_graph_break = compiled_with_graph_break(q, k, v)
    diff = (out_no_comp - out_compiled_with_graph_break).abs().max()
    print(diff) # No real diff

    #### Doesn't work
    compiled = torch.compile(attention_function)
    out_compiled = compiled(q, k, v)
    diff = (out_no_comp - out_compiled).abs().max()
    print(diff) # Clear diff, 5.6ish

Here if we don't introduce a graph break for the contiguous calls, we don't get the same results as with eager.

[ZhangLirong-amd]

@avjves , Hi, you can try it using latest aiter, we fix it with PR #960

[avjves]

@avjves , Hi, you can try it using latest aiter, we fix it with PR #960

Tested, everything seems to work with the new fix :) Thanks!

[jjuvonen-amd]

I can also confirm that contiguous, FA and compile work on gfx950 with the latest fix. Thanks, @ZhangLirong-amd!
I think this issue can be closed now.