[enflame] fix stride mismatch int flash_attention_forward test

[xingxing588]

detail log:
test_attention_ops.py::attention_ref

   output = torch.einsum("bhts,bshd->bthd", attention_drop, drop_v)
   if query_padding_mask is not None:
        output.masked_fill_((~query_padding_mask)[:, :, None, None], 0.0)
    log:output.shape torch.Size([1, 2, 4, 128])
    log:output.stride (256, 128, 256, 1)
    output = output.contiguous() --need contiguous
    log:output.shape torch.Size([1, 2, 4, 128])
    log:output.stride (1024, 512, 128, 1)
flag_gems/ops/attention.py::flash_attention_forward
    query = transpose(1, 2) got a temp layout
    log:out torch.Size([1, 4, 2, 128])
    log:out.stride (1024, 128, 512, 1)
    log:out.is_contiguous() False
    out = out.contiguous()
    log:out.contiguous() torch.Size([1, 4, 2, 128])
    log:out.contiguous().stride (1024, 128, 512, 1)
    log:out.contiguous().is_contiguous() False

PR Category

Type of Change

Description

Issue

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

[CLAassistant]

All committers have signed the CLA.

[kiddyjinjin]

lgtm

[0x45f]

why using torch.empty instead of torch.empty_like