YangXu1990uiuc opened a new pull request, #20078:
URL: https://github.com/apache/tvm/pull/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::GetAttrOpt<T>()`, which returns `std::nullopt`
for an absent attribute
*or* for one that does not hold the requested type, and used it for `scale`.
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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
