This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 5dbf7e05c9 [Relax][cuDNN] Do not offload causal / non-fp16 attention,
and fix the default softmax scale (#20078)
5dbf7e05c9 is described below
commit 5dbf7e05c97ff1eae71eb74d3341b6fc681340e9
Author: Yang Xu <[email protected]>
AuthorDate: Sun Sep 13 20:45:49 2026 -0700
[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.
---
python/tvm/relax/backend/cuda/cudnn.py | 9 +++++++
python/tvm/relax/backend/patterns.py | 1 +
python/tvm/relax/testing/attention.py | 3 ++-
.../extra/contrib/cudnn/cudnn_json_runtime.cc | 11 +++++++-
src/runtime/extra/contrib/json/json_node.h | 17 ++++++++++++
tests/python/relax/test_codegen_cudnn.py | 31 ++++++++++++++++++++++
6 files changed, 70 insertions(+), 2 deletions(-)
diff --git a/python/tvm/relax/backend/cuda/cudnn.py
b/python/tvm/relax/backend/cuda/cudnn.py
index 2df31ca62f..8036719259 100644
--- a/python/tvm/relax/backend/cuda/cudnn.py
+++ b/python/tvm/relax/backend/cuda/cudnn.py
@@ -86,6 +86,15 @@ def _check_stacked_attention(context: PatternCheckContext,
layout: str) -> bool:
return False
else:
raise NotImplementedError(f"Unsupported layout: {layout}")
+ if not context.annotated_expr["stacked_qkv"].ty.dtype == "float16":
+ return False
+ # The cuDNN runtime builds an unmasked SDPA graph, so offloading a masked
attention would
+ # silently compute bidirectional attention.
+ attention = context.annotated_expr["attention"]
+ if attention.attrs.causal_mask is not None:
+ return False
+ if attention.attrs.window_size is not None:
+ return False
return True
diff --git a/python/tvm/relax/backend/patterns.py
b/python/tvm/relax/backend/patterns.py
index 06011d6f6e..a669545819 100644
--- a/python/tvm/relax/backend/patterns.py
+++ b/python/tvm/relax/backend/patterns.py
@@ -336,6 +336,7 @@ def make_stacked_attention_pattern(start_op: str,
with_bias: bool = False, layou
out = is_op("relax.nn.attention_bias")(query, key, value, bias)
else:
out = is_op("relax.nn.attention")(query, key, value)
+ annotations["attention"] = out
if layout == "SBN3H":
out = is_op("relax.permute_dims")(out)
diff --git a/python/tvm/relax/testing/attention.py
b/python/tvm/relax/testing/attention.py
index 449902425d..3ea7d50228 100644
--- a/python/tvm/relax/testing/attention.py
+++ b/python/tvm/relax/testing/attention.py
@@ -75,6 +75,7 @@ def get_relax_stacked_attention_module(
qk_scale=None,
single_shape=False,
layout="BS3NH",
+ causal_mask=None,
): # pylint: disable=too-many-arguments, too-many-locals, too-many-branches,
invalid-name
# pylint: disable=too-many-statements
"""Get a relax module for stacked attention."""
@@ -139,7 +140,7 @@ def get_relax_stacked_attention_module(
q = R.permute_dims(q, [1, 0, 2, 3])
k = R.permute_dims(k, [1, 0, 2, 3])
v = R.permute_dims(v, [1, 0, 2, 3])
- result = R.emit(R.nn.attention(q, k, v, bias, qk_scale))
+ result = R.emit(R.nn.attention(q, k, v, bias, qk_scale,
causal_mask))
if layout == "SBN3H":
result = R.emit(R.permute_dims(result, [1, 0, 2, 3]))
R.output(result)
diff --git a/src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc
b/src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc
index 5b3756fd79..7246983fcd 100644
--- a/src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc
+++ b/src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc
@@ -215,11 +215,20 @@ class cuDNNJSONRuntime : public JSONRuntimeBase {
} else {
TVM_FFI_THROW(InternalError) << "Unsupported layout: " << layout;
}
+ // A `None` scale is serialized as an empty string, so HasAttrValue (not
HasAttr) decides
+ // whether the default applies. GetAttr throws on a type mismatch, so a
scale we cannot
+ // read is loud rather than silently replaced by the default.
double scale = 1 / std::sqrt(head_size);
- if (node.HasAttr("scale")) {
+ if (node.HasAttrValue("scale")) {
scale = node.GetAttr<double>("scale");
}
+ // The SDPA graph below is built without any mask, so masked attention
must not reach here.
+ TVM_FFI_ICHECK(!node.HasAttrValue("causal_mask"))
+ << "cuDNN attention does not support causal_mask yet";
+ TVM_FFI_ICHECK(!node.HasAttrValue("window_size"))
+ << "cuDNN attention does not support window_size yet";
+
auto runner = tvm::contrib::CuDNNSDPARunner::Create();
runner->Init(batch, seq_len, num_heads, num_kv_heads, head_size,
head_size_v, scale, dtype,
layout);
diff --git a/src/runtime/extra/contrib/json/json_node.h
b/src/runtime/extra/contrib/json/json_node.h
index 40c96d8269..da51bb45c6 100644
--- a/src/runtime/extra/contrib/json/json_node.h
+++ b/src/runtime/extra/contrib/json/json_node.h
@@ -257,6 +257,23 @@ class JSONGraphNode {
return attrs_[key].cast<T>();
}
+ /*!
+ * \brief Check whether an optional attribute carries a value.
+ *
+ * The JSON serializer stores a `None` attribute as an empty string, so
`HasAttr` alone
+ * cannot distinguish an unset attribute from one that was set. Callers pair
this with
+ * `GetAttr<T>`, which throws on a type mismatch rather than silently
falling back.
+ *
+ * \param key The key for lookup.
+ *
+ * \return Whether the attribute is present and is not the empty-string
`None` sentinel.
+ */
+ bool HasAttrValue(const std::string& key) const {
+ if (attrs_.count(key) == 0) return false;
+ auto sentinel = attrs_[key].try_cast<ffi::String>();
+ return !(sentinel.has_value() && sentinel.value().empty());
+ }
+
/*!
* \brief Set an attribute for the node.
*
diff --git a/tests/python/relax/test_codegen_cudnn.py
b/tests/python/relax/test_codegen_cudnn.py
index 419c805fdc..b934cf9967 100644
--- a/tests/python/relax/test_codegen_cudnn.py
+++ b/tests/python/relax/test_codegen_cudnn.py
@@ -299,6 +299,37 @@ def stacked_attention_size(request):
return request.param
+def _is_offloaded_to_cudnn(mod):
+ return any("cudnn" in gv.name_hint for gv, _ in mod.functions_items())
+
+
+def _get_stacked_attention_module(dtype, causal_mask=None):
+ b, s, n, h, h_v = 4, 8, 32, 64, 64
+ qkv = np.random.randn(b, s, n * h * 2 + n * h_v).astype(dtype)
+ return get_relax_stacked_attention_module(
+ qkv, b, s, n, h, h_v, "split", causal_mask=causal_mask
+ )
+
+
+def test_stacked_attention_partition():
+ mod = _get_stacked_attention_module("float16")
+ assert _is_offloaded_to_cudnn(partition_for_cudnn(mod))
+
+
[email protected]("causal_mask", ["TopLeft", "BottomRight"])
+def test_stacked_attention_causal_not_partitioned(causal_mask):
+ # The cuDNN runtime builds an unmasked SDPA graph, offloading here would
silently drop
+ # the causal mask.
+ mod = _get_stacked_attention_module("float16", causal_mask=causal_mask)
+ assert not _is_offloaded_to_cudnn(partition_for_cudnn(mod))
+
+
+def test_stacked_attention_fp32_not_partitioned():
+ # attention.cc only builds a half-precision graph and ICHECKs at module
init otherwise.
+ mod = _get_stacked_attention_module("float32")
+ assert not _is_offloaded_to_cudnn(partition_for_cudnn(mod))
+
+
@pytest.mark.skip(reason="require cudnn frontend")
def test_stacked_attention_split_offload(stacked_attention_size):
b, s, n, (h, h_v), bias_shape, scale, single_shape, layout =
stacked_attention_size
