onnxruntime
dbbdc8a1 - [CUDA] Zero the cuDNN SDPA graph cache key before hashing it (#34081)

Commit
1 day ago
[CUDA] Zero the cuDNN SDPA graph cache key before hashing it (#34081) ### Description `cudnn_sdpa::run()` fills the `mha_graph_cache` key (`GraphParams`) one field at a time, but `BytesHash` hashes all `sizeof(GraphParams)` bytes. The padding bytes are never written, so they hold whatever was on the stack. Two calls with equal fields can then hash differently, miss the cache and build the same cuDNN graph again. This change zeroes the struct first, as `FillPagedGraphParams` already does for the paged cache. On x64, `GraphParams` is 64 bytes, and 9 of them are padding: bytes 34-35, 53-55 and 60-63 (printed with `offsetof` under MSVC 19.51 and GCC 13.3). Every `cudnn_sdpa::run` caller shares this cache: contrib MHA, contrib GQA and the ONNX `Attention` op. Refs #33995. This only removes the duplicate build. It doesn't change the one build per new length that #33995 asks about. ### Motivation and Context john-rocky found this on 1.31.1 ([comment on #33424](https://github.com/microsoft/onnxruntime/pull/33424#issuecomment-6730947657)). The Qwen3-0.6B fp16 decoder built the cuDNN graph twice for most new prompt lengths, and the two keys differed only in bytes 34-35. Adding this memset to a v1.31.1 source build brought every length down to one build. I reproduced the duplicate on the stock 1.31.0 wheel, with a different GPU and cuDNN version, and checked main. Setup: RTX 5070 Ti Laptop GPU (SM 12.0), Windows 11, driver 621.01. The model is Qwen3-0.6B exported with the onnxruntime-genai 0.19.0 builder (`-p fp16 -e cuda`), run with the default `sdpa_kernel`, prefill only with no past. It ran twice at 600 tokens, then three times at each other length. Builds are counted as `CUDNN_BACKEND_OPERATIONGRAPH_DESCRIPTOR` finalizes in the cuDNN API log (`CUDNN_LOGLEVEL_DBG=3`): | Graph builds per new length | 600 | 512 | 400 | 300 | 256 | 200 | 150 | 128 | 100 | |---|---|---|---|---|---|---|---|---|---| | onnxruntime-gpu 1.31.0 from PyPI, cuDNN 9.27 | 3 | 1 | 2 | 2 | 2 | 2 | 2 | 2 | 2 | | main from 11-06 (with #33945) built locally, cuDNN 9.20 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | | same build, with this change | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | 1 | The 1.31.0 row matches john-rocky's PyPI 1.31.1 row. My main build doesn't show the duplicate with this model, and a temporary stderr line per cache lookup, added to that build without this change, shows why. All 728 lookups (28 layers × 26 runs) had the same padding: `00 00` at 34-35, `00 00 00` at 53-55 and `0b 00 00 00` at 60-63. Byte 60 is left over from the stack. In that build it happens to be the same on every call, so the lookups still hit. Whether they hit depends on what earlier calls leave at that address, and that differs between the 1.31.x builds and this one. `cudnn_flash_attention.cc` in that build is identical to current main. ### Testing On the main build above, with this change: - `onnxruntime_provider_test`: `GroupQueryAttention*:PagedAttention*` 183 passed / 37 skipped, `MultiHeadAttentionTest*` 33 passed, `AttentionTest.*` 89 passed. With `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1`, cuDNN SDPA was picked 21, 31 and 3 times in those runs. - The Qwen3-0.6B decoder above: one build per new length. - clang-format 20.1.8: no changes. I didn't add a unit test. The cache is local to this file, and whether the padding differs depends on the stack, so a test can't reliably trigger the miss through the public API. The paged key's memset has no dedicated test either. I can add one through the CUDA test bridge if you'd like. #33424 moves this fill into a `FillGraphParams` helper that also zeroes the struct. Whichever of the two merges second will need a one-hunk rebase. About the question on #33995 whether 93e7d33326 (#33123) solves the issue: that commit caches Conv execution plans in `nn/conv.cc` and doesn't touch `cudnn_flash_attention.cc`, so it doesn't change this cache. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Author
Parents
Loading