Sdpa api

[drisspg]

trafficstars

Summary

This gunna take a second

[drisspg]

Failing with

msg = "torch.* op returned non-Tensor _SDPAParams call_function <class 'torch._C._SDPAParams'>"

    def unimplemented(msg: str):
        assert msg != os.environ.get("BREAK", False)
>       raise Unsupported(msg)
E       torch._dynamo.exc.Unsupported: torch.* op returned non-Tensor _SDPAParams call_function <class 'torch._C._SDPAParams'>
E
E       from user code:
E          File "/home/drisspg/meta/transformer_nuggets/transformer_nuggets/sdpa/sdpa.py", line 36, in sdpa_prototype
E           return attn_mask.dispatch(query, key, value, attn_mask, causal, scale, dropout_p)
E         File "/home/drisspg/meta/transformer_nuggets/transformer_nuggets/sdpa/attn_bias.py", line 172, in dispatch
E           sdpa_params = SDPAParams(query, key, value, None, dropout_p, causal)
E
E       Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E
E
E       You can suppress this exception and fall back to eager by setting:
E           import torch._dynamo
E           torch._dynamo.config.suppress_errors = True

../pytorch/torch/_dynamo/exc.py:184: Unsupported