CUDNN: multi-head attention broken on 11.x

[maleadt]

https://buildkite.com/julialang/cuda-dot-jl/builds/5412#019205a9-8db7-4962-b141-2413e84ed18b

multiheadattn: Error During Test at /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/runtests.jl:23
Got exception outside of a @test
LoadError: CUDNNError: CUDNN_STATUS_INTERNAL_ERROR (code 4000)
Stacktrace:
  [1] throw_api_error(res::cuDNN.cudnnStatus_t)
    @ cuDNN /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/libcudnn.jl:14
  [2] check
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/libcudnn.jl:27 [inlined]
  [3] cudnnMultiHeadAttnForward
    @ ~/.cache/julia-buildkite-plugin/depots/3cc01fab-3357-4a7a-9294-cde2d3115a97/packages/CUDA/z3j2H/lib/utils/call.jl:34 [inlined]
  [4] cudnnMultiHeadAttnForwardAD(weights::CuArray{Float32, 1, CUDA.DeviceMemory}, queries::CuArray{Float32, 3, CUDA.DeviceMemory}, keys::CuArray{Float32, 3, CUDA.DeviceMemory}, values::CuArray{Float32, 3, CUDA.DeviceMemory}, residuals::Nothing; dready::Base.RefValue{Bool}, dweights::Base.RefValue{Any}, dqueries::Base.RefValue{Any}, dkeys::Base.RefValue{Any}, dvalues::Base.RefValue{Any}, attnDesc::cudnnAttnDescriptor, currIdx::Int64, loWinIdx::Vector{Int32}, hiWinIdx::Vector{Int32}, devSeqLengthsQO::CuArray{Int32, 1, CUDA.DeviceMemory}, devSeqLengthsKV::CuArray{Int32, 1, CUDA.DeviceMemory}, qDesc::cudnnSeqDataDescriptor, kDesc::cudnnSeqDataDescriptor, vDesc::cudnnSeqDataDescriptor, oDesc::cudnnSeqDataDescriptor, out::CuArray{Float32, 3, CUDA.DeviceMemory}, workspace::Nothing, reserveSpace::Nothing)
    @ cuDNN /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/multiheadattn.jl:145
  [5] cudnnMultiHeadAttnForwardAD
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/multiheadattn.jl:129 [inlined]
  [6] cudnnMultiHeadAttnForward!(out::CuArray{Float32, 3, CUDA.DeviceMemory}, weights::CuArray{Float32, 1, CUDA.DeviceMemory}, queries::CuArray{Float32, 3, CUDA.DeviceMemory}, keys::CuArray{Float32, 3, CUDA.DeviceMemory}, values::CuArray{Float32, 3, CUDA.DeviceMemory}, attnDesc::cudnnAttnDescriptor; axes::Vector{cudnnSeqDataAxis_t}, seqLengthsQO::Vector{Int32}, seqLengthsKV::Vector{Int32}, devSeqLengthsQO::CuArray{Int32, 1, CUDA.DeviceMemory}, devSeqLengthsKV::CuArray{Int32, 1, CUDA.DeviceMemory}, qDesc::cudnnSeqDataDescriptor, kDesc::cudnnSeqDataDescriptor, vDesc::cudnnSeqDataDescriptor, oDesc::cudnnSeqDataDescriptor, residuals::Nothing, currIdx::Int64, loWinIdx::Vector{Int32}, hiWinIdx::Vector{Int32}, workspace::Nothing, reserveSpace::Nothing, dweights::Base.RefValue{Any}, dqueries::Base.RefValue{Any}, dkeys::Base.RefValue{Any}, dvalues::Base.RefValue{Any}, o::@Kwargs{attnMode::UInt32, nHeads::Int32, smScaler::Float64, mathType::cudnnMathType_t, qProjSize::Int64, kProjSize::Int64, vProjSize::Int64, oProjSize::Int64, qoMaxSeqLength::Int64, kvMaxSeqLength::Int64, maxBatchSize::Int64, maxBeamSize::Int64})
    @ cuDNN /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/multiheadattn.jl:118
  [7] cudnnMultiHeadAttnForward
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/multiheadattn.jl:55 [inlined]
  [8] cudnnMultiHeadAttnForward(w::CuArray{Float32, 1, CUDA.DeviceMemory}, q::CuArray{Float32, 3, CUDA.DeviceMemory}, k::CuArray{Float32, 3, CUDA.DeviceMemory}, v::CuArray{Float32, 3, CUDA.DeviceMemory}; o::@Kwargs{axes::Vector{cudnnSeqDataAxis_t}, seqLengthsQO::Vector{Int32}, seqLengthsKV::Vector{Int32}, attnMode::UInt32, nHeads::Int32, smScaler::Float64, mathType::cudnnMathType_t, qProjSize::Int64, kProjSize::Int64, vProjSize::Int64, oProjSize::Int64, qoMaxSeqLength::Int64, kvMaxSeqLength::Int64, maxBatchSize::Int64, maxBeamSize::Int64, residuals::Nothing, currIdx::Int64, loWinIdx::Vector{Int32}, hiWinIdx::Vector{Int32}})
    @ cuDNN /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/src/multiheadattn.jl:54
  [9] mhatest(; axes::Vector{cudnnSeqDataAxis_t}, seqLengthsQO::Vector{Int32}, seqLengthsKV::Vector{Int32}, attnMode::UInt32, nHeads::Int32, smScaler::Float64, mathType::cudnnMathType_t, qProjSize::Int64, kProjSize::Int64, vProjSize::Int64, oProjSize::Int64, qoMaxSeqLength::Int64, kvMaxSeqLength::Int64, maxBatchSize::Int64, maxBeamSize::Int64, residuals::Nothing, currIdx::Int64, loWinIdx::Vector{Int32}, hiWinIdx::Vector{Int32})
    @ Main /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/multiheadattn.jl:113
 [10] mhatest()
    @ Main /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/multiheadattn.jl:55
 [11] top-level scope
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/multiheadattn.jl:135
 [12] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [13] macro expansion
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/runtests.jl:24 [inlined]
 [14] macro expansion
    @ ~/.cache/julia-buildkite-plugin/julia_installs/bin/linux/x64/1.10/julia-1.10-latest-linux-x86_64/share/julia/stdlib/v1.10/Test/src/Test.jl:1577 [inlined]
 [15] top-level scope
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/runtests.jl:24
 [16] eval
    @ ./boot.jl:385 [inlined]
 [17] macro expansion
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/runtests.jl:22 [inlined]
 [18] macro expansion
    @ ~/.cache/julia-buildkite-plugin/julia_installs/bin/linux/x64/1.10/julia-1.10-latest-linux-x86_64/share/julia/stdlib/v1.10/Test/src/Test.jl:1577 [inlined]
 [19] top-level scope
    @ /var/lib/buildkite-agent/builds/gpuci-16/julialang/cuda-dot-jl/lib/cudnn/test/runtests.jl:13
 [20] include(fname::String)
    @ Base.MainInclude ./client.jl:489
 [21] top-level scope
    @ none:6
 [22] eval
    @ ./boot.jl:385 [inlined]
 [23] exec_options(opts::Base.JLOptions)
    @ Base ./client.jl:291
 [24] _start()
    @ Base ./client.jl:552