The GitHub Actions job "Lint" on tvm.git/main has succeeded. Run started by GitHub user tlopex (triggered by tlopex).
Head commit for run: 5dbf7e05c97ff1eae71eb74d3341b6fc681340e9 / Yang Xu <[email protected]> [Relax][cuDNN] Do not offload causal / non-fp16 attention, and fix the default softmax scale (#20078) # [Relax][cuDNN] Do not offload causal / non-fp16 attention, and fix the default softmax scale ## Summary The cuDNN BYOC backend offloads any `relax.nn.attention` that matches the stacked-QKV pattern, without ever looking at the op's attributes. A model that asks for causal attention is therefore compiled into a cuDNN SDPA graph that computes **full bidirectional attention** — silently, with no error and no warning. The same blind spot lets fp32/bf16 attention be partitioned to a runtime that only builds a half-precision graph, and makes the documented `1/sqrt(head_dim)` default softmax scale unreachable. This PR rejects the workloads cuDNN is not actually being asked to compute, fixes the default-scale path, and adds partition-level regression tests. ### 1. Causal attention is offloaded and computed non-causally `make_stacked_attention_pattern` (python/tvm/relax/backend/patterns.py:338) matches the attention op node only, with no constraint on its attributes, and `_check_stacked_attention` (python/tvm/relax/backend/cuda/cudnn.py:69-89) validates only ndim and the split axis. The runtime then builds the SDPA graph with `.set_causal_mask(false)` (src/runtime/extra/contrib/cudnn/cudnn_frontend/attention.cc:88-93). On the cuDNN side that call is a no-op: in the cuDNN frontend, `SDPA_attributes::set_causal_mask(bool value)` is guarded by `if (value)` and only ever *adds* `set_diagonal_alignment(TOP_LEFT)` + `set_diagonal_band_right_bound(0)` (graph_properties.h, `set_causal_mask`). Passing `false` leaves the graph in its default state — no diagonal band bound at all, i.e. every query attends to every key. Failure scenario: `R.nn.attention(q, k, v, causal_mask="TopLeft")` inside the stacked-QKV pattern is partitioned into `fused_..._cudnn`, and the compiled model returns bidirectional attention output. Decoder models produce wrong (future-leaking) results with no diagnostic. Fix (conservative option): reject the offload at partition time. `_check_stacked_attention` now returns `False` when the attention op carries `causal_mask` or `window_size`. To make the op's attributes visible to the check, the attention call is added to the pattern annotations (`annotations["attention"]`), the same idiom already used for `split` / `q_transpose` and for `root` in `make_conv2d_pattern`. A defensive `ICHECK` was also added in the JSON runtime so that any future pattern cannot re-introduce the silent-wrong-answer path. I deliberately did **not** plumb the attribute through to `attention.cc`. Doing it correctly means mapping `"TopLeft"`/`"BottomRight"` onto `set_diagonal_alignment(DiagonalAlignment_t::TOP_LEFT / BOTTOM_RIGHT)` + `set_diagonal_band_right_bound(0)`, and — for `window_size` — onto `set_diagonal_band_left_bound()`; note `set_sliding_window_length()` in the cuDNN frontend is a pure alias for `set_diagonal_band_left_bound()` and sets only the *left* bound, so a sliding-window implementation must set both bounds explicitly. That is a feature addition that needs numerical validation on a GPU, which this audit could not run. Rejecting the pattern falls back to the non-offloaded path, which is correct. ### 2. The default softmax scale never reaches the runtime `AttentionAttrs::scale` is `Optional<FloatImm>`. When it is `None` the JSON serializer writes an empty **string** for the attribute (src/relax/backend/contrib/codegen_json/codegen_json.h:188-191, and the `Optional` overloads at :76-90). The runtime then does (src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc:218-221): ```cpp double scale = 1 / std::sqrt(head_size); if (node.HasAttr("scale")) { scale = node.GetAttr<double>("scale"); } ``` `HasAttr("scale")` is therefore *always* true — the `1/sqrt(head_size)` default is dead code — and `GetAttr<double>` (src/runtime/extra/contrib/json/json_node.h:254-258) casts an `ffi::String` to `double`, which fails. Either way the documented default scale never reaches cuDNN. Failure scenario: any `R.nn.attention(q, k, v)` written without an explicit `scale` — the common case, and two of the four parametrizations of the existing `test_stacked_attention_split_offload`. Fix: added `JSONGraphNode::HasAttrValue()`, which reports whether an attribute is present and is not the empty-string `None` sentinel, and paired it with the existing `GetAttr<T>()` for `scale`. `GetAttr<T>` throws on a type mismatch, so an attribute we cannot read is loud rather than silently replaced by the default -- which would be the same class of silent wrong answer this PR exists to fix. This keeps the "`None` is an empty string" serialization convention that every other BYOC runtime already relies on, instead of changing the shared serializer. ### 3. No dtype guard at partition time `_check_stacked_attention` (python/tvm/relax/backend/cuda/cudnn.py:69) has no dtype check, but the runtime only ever builds a half-precision graph and hard-fails on `TVM_FFI_ICHECK(data_type.code == kDLFloat && data_type.bits == 16) << "Only float16 is supported"` (src/runtime/extra/contrib/cudnn/cudnn_frontend/attention.cc:42-43). Failure scenario: an fp32 (or bf16) attention module is happily partitioned to cuDNN and then aborts at module init with "Only float16 is supported", instead of falling back to a working non-offloaded implementation. Fix: require `float16` on the stacked QKV input, following the `_is_supported_dtype` idiom already used by `_check_conv2d`. ## Evidence Confirmed by reading the full chain (pattern -> offload check -> runtime). No runtime repro was performed (TVM was not built), so no measured numbers are claimed for this repo. ## Testing Added to `tests/python/relax/test_codegen_cudnn.py`: * `test_stacked_attention_partition` — positive control: fp16 unmasked attention is still partitioned to cuDNN (guards against over-rejection). * `test_stacked_attention_causal_not_partitioned[TopLeft|BottomRight]` — fails before this change (the causal graph is offloaded), passes after. * `test_stacked_attention_fp32_not_partitioned` — fails before this change, passes after. These are partition-only tests and need no GPU, but note that `test_codegen_cudnn.py` sets a module-level `pytestmark = [pytest.mark.gpu, skipif(not env.has_cudnn())]`, so they are skipped in environments without cuDNN, as is the existing `test_cudnn_partition_conv2d_without_bias`. `get_relax_stacked_attention_module` gained an optional `causal_mask` argument to build the negative cases. **None of this was run locally**: TVM is not built in the audit environment, so the new tests are unverified. The changed Python files were byte-compiled and checked with the repo's formatting settings (100-column), and the changed C++ files are clean under the repo `.clang-format`. Also worth flagging for maintainers: the only end-to-end cuDNN SDPA test, `test_stacked_attention_split_offload`, is unconditionally skipped (`@pytest.mark.skip(reason="require cudnn frontend")`, tests/python/relax/test_codegen_cudnn.py:302). Nothing in CI exercises this runtime, which is how findings 1 and 2 survived. Found by an integration audit of cuDNN SDPA consumers by the NVIDIA cuDNN team. Report URL: https://github.com/apache/tvm/actions/runs/34803705967 With regards, GitHub Actions via GitBox --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
