[Relax][cuDNN] Do not offload causal / non-fp16 attention, and fix the default softmax scale

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
 
 

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)

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)

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);

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.
    *

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