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

[YangXu1990uiuc]

[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):

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.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.