From d76e4c620672928dfce7c762c739e0ff37b49d7b Mon Sep 17 00:00:00 2001 From: Yang Xu Date: Thu, 30 Jul 2026 01:25:12 -0700 Subject: [PATCH] [Relax][cuDNN] Reject unsupported attention offloads; fix default scale The cuDNN BYOC backend matched relax.nn.attention on the op node alone, so attention carrying a causal mask was offloaded to a cuDNN SDPA graph built with .set_causal_mask(false) and silently computed bidirectionally. - python/tvm/relax/backend/patterns.py:339: annotate the attention call so the partition check can read the op attributes. - python/tvm/relax/backend/cuda/cudnn.py:89-97: do not offload when causal_mask or window_size is set, or when the input is not float16 (the runtime only builds a half-precision graph, attention.cc:42). - src/runtime/extra/contrib/json/json_node.h:261-274: add HasAttrValue, since a None attribute is serialized as an empty string and HasAttr cannot tell. It is paired with GetAttr, which throws on a type mismatch, so an attribute we cannot read is loud rather than silently replaced by the default. - src/runtime/extra/contrib/cudnn/cudnn_json_runtime.cc:218-231: use it so the 1/sqrt(head_size) default scale is actually applied, and reject masked attention that reaches the runtime. - python/tvm/relax/testing/attention.py:79: optional causal_mask argument. - tests/python/relax/test_codegen_cudnn.py:302-329: partition tests for the causal and fp32 cases, plus a positive control. --- 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 2df31ca62f5d..803671925906 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 06011d6f6e97..a6695458195a 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 449902425dd0..3ea7d50228df 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 5b3756fd79b8..7246983fcd76 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("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 40c96d826914..da51bb45c6ce 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(); } + /*! + * \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`, 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(); + 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 419c805fdc49..b934cf99674e 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)) + + +@pytest.mark.parametrize("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