Fused attention error while running Nvidia Cosmos

[deepbeepmeep]

Hello

I am trying to run the latest Nvidia Cosmos model on a RTX 4090 and I get an error when fused attention is called :
Line 1080 in fused_attn.py / fused_attn_forward /output_tensors = tex.fused_attn_fwd(...)

The transformer_engine compilation didn't produce any error during the installation and I have CuDNN v9.6.0 installed. I have Flash Attention 2.7.3, could this be the issue (there is a warning that says that 2.6.3. is the latest supported) ? is Flash attention used behind the scene ?

E! CuDNN (v90100 70) function cudnnBackendFinalize() called:
e! Error: CUDNN_STATUS_EXECUTION_FAILED; Reason: rtc->loadModule()
e! Error: CUDNN_STATUS_EXECUTION_FAILED; Reason: ptr.isSupported()
e! Error: CUDNN_STATUS_EXECUTION_FAILED; Reason: engine_post_checks(*engine_iface, engine.getPerfKnobs(), req_size, engine.getTargetSMCount())
e! Error: CUDNN_STATUS_EXECUTION_FAILED; Reason: finalize_internal()
e! Time: 2025-01-14T00:21:03.310308 (0d+0h+3m+34s since start)
e! Process=381629; Thread=381629; GPU=NULL; Handle=NULL; StreamId=NULL.

Many thanks in davance